chore: import upstream snapshot with attribution

tests/python/json_ffi/test_json_ffi_engine.py

@@ -0,0 +1,197 @@
+import json
+from typing import Dict, List, Optional  # noqa: UP035
+
+import pytest
+from pydantic import BaseModel
+
+from mlc_llm.json_ffi import JSONFFIEngine
+from mlc_llm.testing import require_test_model
+
+# test category "engine_feature"
+pytestmark = [pytest.mark.engine_feature]
+
+
+chat_completion_prompts = [
+    "What is the meaning of life?",
+    "Introduce the history of Pittsburgh to me. Please elaborate in detail.",
+    "Write a three-day Seattle travel plan. Please elaborate in detail.",
+    "What is Alaska famous of? Please elaborate in detail.",
+    "What is the difference between Lambda calculus and Turing machine? Please elaborate in detail.",  # noqa: E501
+    "What are the necessary components to assemble a desktop computer? Please elaborate in detail.",
+    "Why is Vitamin D important to human beings? Please elaborate in detail.",
+    "Where is milk tea originated from? Please elaborate in detail.",
+    "Where is the southernmost place in United States? Please elaborate in detail.",
+    "Do you know AlphaGo? What capabilities does it have, and what achievements has it got? Please elaborate in detail.",  # noqa: E501
+]
+
+function_calling_prompts = [
+    "What is the temperature in Pittsburgh, PA?",
+    "What is the temperature in Tokyo, JP?",
+    "What is the temperature in Pittsburgh, PA and Tokyo, JP?",
+]
+
+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"],
+            },
+        },
+    }
+]
+
+
+def run_chat_completion(
+    engine: JSONFFIEngine,
+    model: str,
+    prompts: List[str] = chat_completion_prompts,  # noqa: UP006
+    tools: Optional[List[Dict]] = None,  # noqa: UP006
+):
+    num_requests = 2
+    max_tokens = 64
+    n = 1
+    output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)]  # noqa: UP006
+
+    for rid in range(num_requests):
+        print(f"chat completion for request {rid}")
+        for response in engine.chat.completions.create(
+            messages=[{"role": "user", "content": [{"type": "text", "text": prompts[rid]}]}],
+            model=model,
+            max_tokens=max_tokens,
+            n=n,
+            request_id=str(rid),
+            tools=tools,
+        ):
+            for choice in response.choices:
+                assert choice.delta.role == "assistant"
+                assert isinstance(choice.delta.content, str)
+                output_texts[rid][choice.index] += choice.delta.content
+
+    # Print output.
+    print("Chat completion all finished")
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+
+def run_json_schema_function_calling(
+    engine: JSONFFIEngine,
+    model: str,
+    prompts: List[str] = function_calling_prompts,  # noqa: UP006
+    tools: Optional[List[Dict]] = None,  # noqa: UP006
+):
+    num_requests = 2
+    max_tokens = 64
+    n = 1
+    output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)]  # noqa: UP006
+
+    class ToolCall(BaseModel):
+        name: str
+        arguments: Dict[str, str]  # noqa: UP006
+
+    class Schema(BaseModel):
+        tool_calls: List[ToolCall]  # noqa: UP006
+
+    schema_str = json.dumps(Schema.model_json_schema())
+    print("Schema str", schema_str)
+
+    for rid in range(num_requests):
+        print(f"chat completion for request {rid}")
+        for response in engine.chat.completions.create(
+            messages=[
+                {
+                    "role": "system",
+                    "content": "You are a function calling AI model. You are provided with function signatures within "  # noqa: E501
+                    "<tools></tools> XML tags. You may call one or more functions to assist with the user query. Don't make "  # noqa: E501
+                    f"assumptions about what values to plug into functions. Here are the available tools: <tools> {json.dumps(tools)} </tools> "  # noqa: E501
+                    "Do not stop calling functions until the task has been accomplished or you've reached max iteration of 10. "  # noqa: E501
+                    "Calling multiple functions at once can overload the system and increase cost so call one function at a time please. "  # noqa: E501
+                    "If you plan to continue with analysis, always call another function. Return a valid json object (using double "  # noqa: E501
+                    f"quotes) in the following schema: {schema_str}",
+                },
+                {"role": "user", "content": [{"type": "text", "text": prompts[rid]}]},
+            ],
+            model=model,
+            max_tokens=max_tokens,
+            n=n,
+            request_id=str(rid),
+            response_format={"type": "json_object", "schema": schema_str},
+        ):
+            for choice in response.choices:
+                assert choice.delta.role == "assistant"
+                assert isinstance(choice.delta.content, str)
+                output_texts[rid][choice.index] += choice.delta.content
+
+    # Print output.
+    print("Chat completion all finished")
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+
+@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC")
+def test_chat_completion(model):
+    # Create engine.
+    engine = JSONFFIEngine(model)
+
+    run_chat_completion(engine, model)
+
+    # Test malformed requests.
+    for response in engine._raw_chat_completion(
+        "malformed_string", include_usage=False, request_id="123"
+    ):
+        assert len(response.choices) == 1
+        assert response.choices[0].finish_reason == "error"
+
+    engine.terminate()
+
+
+@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC")
+def test_reload_reset_unload(model):
+    # Create engine.
+    engine = JSONFFIEngine(model)
+
+    # Run chat completion before and after reload/reset.
+    run_chat_completion(engine, model)
+    engine._test_reload()
+    run_chat_completion(engine, model)
+    engine._test_reset()
+    run_chat_completion(engine, model)
+    engine._test_unload()
+
+    engine.terminate()
+
+
+@require_test_model("Hermes-2-Pro-Mistral-7B-q4f16_1-MLC")
+def test_json_schema_with_system_prompt(model):
+    engine = JSONFFIEngine(model)
+
+    # run function calling
+    run_json_schema_function_calling(engine, model, function_calling_prompts, tools)
+
+    engine.terminate()
+
+
+if __name__ == "__main__":
+    test_chat_completion()
+    test_reload_reset_unload()
+    test_json_schema_with_system_prompt()

tests/python/json_ffi/test_json_ffi_engine_image.py

@@ -0,0 +1,92 @@
+import base64
+from typing import Dict, List, Optional  # noqa: UP035
+
+import requests
+
+from mlc_llm.json_ffi import JSONFFIEngine
+from mlc_llm.testing import require_test_model
+
+
+def base64_encode_image(url: str) -> str:
+    response = requests.get(url)
+    response.raise_for_status()  # Ensure we got a successful response
+    image_data = base64.b64encode(response.content)
+    image_data_str = image_data.decode("utf-8")
+    data_url = f"data:image/jpeg;base64,{image_data_str}"
+    return data_url
+
+
+image_prompts = [
+    [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image_url",
+                    "image_url": f"{base64_encode_image('https://llava-vl.github.io/static/images/view.jpg')}",
+                },
+                {"type": "text", "text": "What does the image represent?"},
+            ],
+        }
+    ]
+]
+
+
+def run_chat_completion(
+    engine: JSONFFIEngine,
+    model: str,
+    prompts: List[List[Dict]] = image_prompts,  # noqa: UP006
+    tools: Optional[List[Dict]] = None,  # noqa: UP006
+):
+    num_requests = 1
+    max_tokens = 64
+    n = 1
+    output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)]  # noqa: UP006
+
+    for rid in range(num_requests):
+        print(f"chat completion for request {rid}")
+        for response in engine.chat.completions.create(
+            messages=prompts[rid],
+            model=model,
+            max_tokens=max_tokens,
+            n=n,
+            request_id=str(rid),
+            tools=tools,
+        ):
+            for choice in response.choices:
+                assert choice.delta.role == "assistant"
+                assert isinstance(choice.delta.content[0], Dict)  # noqa: UP006
+                assert choice.delta.content[0]["type"] == "text"
+                output_texts[rid][choice.index] += choice.delta.content[0]["text"]
+
+    # Print output.
+    print("Chat completion all finished")
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+
+@require_test_model("llava-1.5-7b-hf-q4f16_1-MLC")
+def test_chat_completion():
+    # Create engine.
+    engine = JSONFFIEngine(
+        model,  # noqa: F821
+        max_total_sequence_length=1024,
+    )
+
+    run_chat_completion(engine, model)  # noqa: F821
+
+    # Test malformed requests.
+    for response in engine._raw_chat_completion("malformed_string", n=1, request_id="123"):
+        assert len(response.choices) == 1
+        assert response.choices[0].finish_reason == "error"
+
+    engine.terminate()
+
+
+if __name__ == "__main__":
+    test_chat_completion()

tests/python/json_ffi/test_json_ffi_engine_mock.py

@@ -0,0 +1,105 @@
+import json
+
+import pytest
+import tvm
+
+from mlc_llm.json_ffi import JSONFFIEngine
+from mlc_llm.testing import require_test_model
+
+# test category "unittest"
+pytestmark = [pytest.mark.unittest]
+
+
+def check_error_handling(engine, expect_str, **params):
+    """Check error handling in raw completion API"""
+    body = {
+        "messages": [{"role": "user", "content": "hello"}],
+        "stream_options": {"include_usage": True},
+    }
+    body.update(params)
+
+    for response in engine._raw_chat_completion(
+        json.dumps(body), include_usage=False, request_id="123"
+    ):
+        if response.choices[0].finish_reason is not None:
+            break
+    if response.choices[0].finish_reason != "error":
+        raise RuntimeError(f"expect the request {params} to hit an error")
+
+    if expect_str not in response.choices[0].delta.content:
+        raise RuntimeError(
+            f"expect '{expect_str}' in error msg, but get '{response.choices[0].delta.content}'"
+        )
+
+
+# NOTE: we only need tokenizers in folder
+# launch time of mock test is fast so we can put it in unittest
+@require_test_model("Llama-3-8B-Instruct-q4f16_1-MLC")
+def test_chat_completion_misuse(model: str):
+    engine = JSONFFIEngine(model, tvm.cpu(), model_lib="mock://echo")
+    # Test malformed requests.
+    for response in engine._raw_chat_completion(
+        "malformed_string", include_usage=False, request_id="123"
+    ):
+        assert len(response.choices) == 1
+        assert response.choices[0].finish_reason == "error"
+    # check parameters
+    check_error_handling(engine, "should be non-negative", temperature=-1)
+    check_error_handling(engine, "in range [0, 1]", top_p=100)
+    check_error_handling(engine, "frequency_penalty", frequency_penalty=100)
+
+
+def check_normal_param_passing(engine):
+    json_schema = """
+    {"properties": {"result": {"items": {"type": "Integer"}, "title": "Result", "type": "array"}},
+      "required": ["result"], "title": "Output", "type": "object"}
+    """
+    param_dict = {
+        "top_p": 0.6,
+        "temperature": 0.8,
+        "frequency_penalty": 0.1,
+        "presence_penalty": 0.1,
+    }
+    usage = None
+    for response in engine.chat.completions.create(
+        messages=[{"role": "user", "content": "hello"}],
+        stream=True,
+        stream_options={"include_usage": True},
+        response_format={"type": "json_object", "schema": json_schema},
+        **param_dict,
+    ):
+        if response.usage is not None:
+            usage = response.usage
+
+    # echo mock will echo back the generation config
+    for k, v in param_dict.items():
+        assert usage.extra[k] == v, f"{k} mismatch"
+    assert "response_format" in usage.extra
+    assert usage.extra["response_format"]["type"] == "json_object"
+    assert "schema" in usage.extra["response_format"]
+
+
+def check_n_generation(engine):
+    hit_set = set()
+    for response in engine.chat.completions.create(
+        messages=[{"role": "user", "content": "hello"}],
+        stream=True,
+        stream_options={"include_usage": True},
+        n=3,
+    ):
+        for choice in response.choices:
+            hit_set.add(choice.index)
+    for i in range(3):
+        assert i in hit_set, f"{i} not in n generation"
+
+
+@require_test_model("Llama-3-8B-Instruct-q4f16_1-MLC")
+def test_chat_completion_api(model: str):
+    engine = JSONFFIEngine(model, tvm.cpu(), model_lib="mock://echo")
+    check_normal_param_passing(engine)
+    check_n_generation(engine)
+
+
+if __name__ == "__main__":
+    test_chat_completion_api()
+    test_chat_completion_misuse()