chore: import upstream snapshot with attribution

tests/python/serve/server/test_server_function_call.py

@@ -0,0 +1,208 @@
+"""
+Test script for function call in chat completion. To run this script, use the following command:
+MLC_SERVE_MODEL_LIB=dist/gorilla-openfunctions-v1-q4f16_1_MLC/gorilla-openfunctions-v1-q4f16_1-cuda.so
+MLC_SERVE_MODEL_LIB=${MLC_SERVE_MODEL_LIB} python -m pytest -x tests/python/serve/server/test_server_function_call.py
+"""  # noqa: E501
+
+import json
+import os
+from typing import Dict, List, Optional, Tuple  # noqa: UP035
+
+import pytest
+import requests
+
+OPENAI_V1_CHAT_COMPLETION_URL = "http://127.0.0.1:8000/v1/chat/completions"
+
+
+def check_openai_nonstream_response(
+    response: Dict,  # noqa: UP006
+    *,
+    model: str,
+    object_str: str,
+    num_choices: int,
+    finish_reason: List[str],  # noqa: UP006
+    completion_tokens: Optional[int] = None,
+):
+    print(response)
+    assert response["model"] == model
+    assert response["object"] == object_str
+
+    choices = response["choices"]
+    assert isinstance(choices, list)
+    assert len(choices) == num_choices
+    for idx, choice in enumerate(choices):
+        assert choice["index"] == idx
+        assert choice["finish_reason"] in finish_reason
+
+        # text: str
+        message = choice["message"]
+        assert message["role"] == "assistant"
+        if choice["finish_reason"] == "tool_calls":
+            assert message["content"] is None
+            assert isinstance(message["tool_calls"], list)
+        else:
+            assert message["tool_calls"] is None
+            assert message["content"] is not None
+
+    usage = response["usage"]
+    assert isinstance(usage, dict)
+    assert usage["total_tokens"] == usage["prompt_tokens"] + usage["completion_tokens"]
+    assert usage["prompt_tokens"] > 0
+
+    if completion_tokens is not None:
+        assert usage["completion_tokens"] == completion_tokens
+
+
+def check_openai_stream_response(
+    responses: List[Dict],  # noqa: UP006
+    *,
+    model: str,
+    object_str: str,
+    num_choices: int,
+    finish_reason: str,
+    echo_prompt: Optional[str] = None,
+    suffix: Optional[str] = None,
+    stop: Optional[List[str]] = None,  # noqa: UP006
+    require_substr: Optional[List[str]] = None,  # noqa: UP006
+):
+    assert len(responses) > 0
+
+    finished = [False for _ in range(num_choices)]
+    outputs = ["" for _ in range(num_choices)]
+    for response in responses:
+        assert response["model"] == model
+        assert response["object"] == object_str
+
+        choices = response["choices"]
+        assert isinstance(choices, list)
+        assert len(choices) == num_choices
+        for idx, choice in enumerate(choices):
+            assert choice["index"] == idx
+
+            delta = choice["delta"]
+            assert delta["role"] == "assistant"
+            assert isinstance(delta["content"], str)
+            outputs[idx] += delta["content"]
+
+            if finished[idx]:
+                assert choice["finish_reason"] == finish_reason
+            elif choice["finish_reason"] is not None:
+                assert choice["finish_reason"] == finish_reason
+                finished[idx] = True
+
+    for output in outputs:
+        if echo_prompt is not None:
+            assert output.startswith(echo_prompt)
+        if suffix is not None:
+            assert output.endswith(suffix)
+        if stop is not None:
+            for stop_str in stop:
+                assert stop_str not in output
+        if require_substr is not None:
+            for substr in require_substr:
+                assert substr in output
+
+
+tools = [
+    {
+        "type": "function",
+        "function": {
+            "name": "get_current_weather",
+            "description": "Get the current weather in a given location",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "location": {
+                        "type": "string",
+                        "description": "The city and state, e.g. San Francisco, CA",
+                    },
+                    "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
+                },
+                "required": ["location"],
+            },
+        },
+    }
+]
+
+
+CHAT_COMPLETION_MESSAGES = [
+    # messages #0
+    [
+        {
+            "role": "user",
+            "content": "What is the current weather in Pittsburgh, PA?",
+        }
+    ],
+    # messages #1
+    [
+        {
+            "role": "user",
+            "content": "What is the current weather in Pittsburgh, PA and Tokyo, JP?",
+        }
+    ],
+    # messages #2
+    [
+        {
+            "role": "user",
+            "content": "What is the current weather in Pittsburgh, PA in fahrenheit?",
+        }
+    ],
+]
+
+
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("messages", CHAT_COMPLETION_MESSAGES)
+def test_openai_v1_chat_completion_function_call(
+    served_model: Tuple[str, str],  # noqa: UP006
+    launch_server,
+    stream: bool,
+    messages: List[Dict[str, str]],  # noqa: UP006
+):
+    # `served_model` and `launch_server` are pytest fixtures
+    # defined in conftest.py.
+
+    payload = {
+        "model": served_model[0],
+        "messages": messages,
+        "stream": stream,
+        "tools": tools,
+    }
+
+    response = requests.post(OPENAI_V1_CHAT_COMPLETION_URL, json=payload, timeout=60)
+    if not stream:
+        check_openai_nonstream_response(
+            response.json(),
+            model=served_model[0],
+            object_str="chat.completion",
+            num_choices=1,
+            finish_reason=["tool_calls", "error"],
+        )
+    else:
+        responses = []
+        for chunk in response.iter_lines(chunk_size=512):
+            if not chunk or chunk == b"data: [DONE]":
+                continue
+            responses.append(json.loads(chunk.decode("utf-8")[6:]))
+        check_openai_stream_response(
+            responses,
+            model=served_model[0],
+            object_str="chat.completion.chunk",
+            num_choices=1,
+            finish_reason="tool_calls",
+        )
+
+
+if __name__ == "__main__":
+    model_lib = os.environ.get("MLC_SERVE_MODEL_LIB")
+    if model_lib is None:
+        raise ValueError(
+            'Environment variable "MLC_SERVE_MODEL_LIB" not found. '
+            "Please set it to model lib compiled by MLC LLM "
+            "(e.g., `./dist/gorilla-openfunctions-v1-q4f16_1_MLC/gorilla-openfunctions-v1-q4f16_1-cuda.so`) "  # noqa: E501
+            "which supports function calls."
+        )
+    MODEL = (os.path.dirname(model_lib), model_lib)
+
+    for msg in CHAT_COMPLETION_MESSAGES:
+        test_openai_v1_chat_completion_function_call(MODEL, None, stream=False, messages=msg)
+        test_openai_v1_chat_completion_function_call(MODEL, None, stream=True, messages=msg)