conftest.py 4.22 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

import pytest


@pytest.fixture
def sample_prompts():
    return [
        "Hello, my name is",
        "The president of the United States is",
        "The capital of France is",
        "The future of AI is",
    ]


@pytest.fixture
def sample_token_ids():
    return [
        [0],
        [0, 1],
        [0, 2, 1],
        [0, 3, 1, 2],
    ]


@pytest.fixture
def sample_regex():
29
30
31
32
    return (
        r"((25[0-5]|(2[0-4]|1\d|[1-9]|)\d)\.){3}"
        r"(25[0-5]|(2[0-4]|1\d|[1-9]|)\d)"
    )
33
34


35
# Note: Ensure this only uses attributes compatible with xgrammar
36
37
38
39
40
@pytest.fixture
def sample_json_schema():
    return {
        "type": "object",
        "properties": {
41
42
            "name": {"type": "string"},
            "age": {"type": "integer"},
43
44
45
46
            "skills": {
                "type": "array",
                "items": {
                    "type": "string",
47
                },
48
            },
49
50
            "grade": {
                "type": "string",
51
                "pattern": "^[A-D]$",  # Regex pattern
52
53
54
            },
            "email": {
                "type": "string",
55
                "pattern": "^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}$",
56
            },
57
58
59
60
61
            "work_history": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
62
                        "company": {"type": "string"},
63
                        "duration": {
64
65
66
                            "type": "number",
                            "minimum": 0.0,
                            "maximum": 100.0,  # Numeric range
67
                        },
68
                        "position": {"type": "string"},
69
                    },
70
                    "required": ["company", "duration", "position"],
71
                    "additionalProperties": False,
72
73
                },
                "minItems": 0,
74
75
                "maxItems": 3,
            },
76
        },
77
78
        "required": ["name", "age", "skills", "grade", "email", "work_history"],
        "additionalProperties": False,
79
80
81
    }


82
# A schema unsupported by xgrammar
83
@pytest.fixture
84
def unsupported_json_schema():
85
86
87
88
89
    return {
        "type": "object",
        "properties": {
            "score": {
                "type": "integer",
90
                "multipleOf": 5,  # Numeric multiple
91
92
93
            },
            "tags": {
                "type": "array",
94
95
                "items": {"type": "string", "minLength": 10, "maxLength": 20},
            },
96
        },
97
        "required": ["score", "tags"],
98
        "additionalProperties": False,
99
100
101
102
103
104
    }


@pytest.fixture
def sample_definition_json_schema():
    return {
105
106
107
108
109
        "$defs": {
            "Step": {
                "properties": {
                    "explanation": {"title": "Explanation", "type": "string"},
                    "output": {"title": "Output", "type": "string"},
110
                },
111
112
113
                "required": ["explanation", "output"],
                "title": "Step",
                "type": "object",
114
115
            }
        },
116
117
118
119
120
        "properties": {
            "steps": {
                "items": {"$ref": "#/$defs/Step"},
                "title": "Steps",
                "type": "array",
121
            },
122
            "final_answer": {"title": "Final Answer", "type": "string"},
123
        },
124
125
126
127
        "required": ["steps", "final_answer"],
        "title": "MathReasoning",
        "type": "object",
        "additionalProperties": False,
128
129
130
131
    }


@pytest.fixture
132
def sample_structured_outputs_choices():
133
    return [
134
135
136
137
138
139
140
141
142
143
        "Python",
        "Java",
        "JavaScript",
        "C++",
        "C#",
        "PHP",
        "TypeScript",
        "Ruby",
        "Swift",
        "Kotlin",
144
145
146
147
    ]


@pytest.fixture
148
149
150
151
152
153
154
155
156
157
158
159
160
def sample_sql_ebnf():
    return """
root ::= select_statement
select_statement ::= "SELECT" column "from" table "where" condition
column ::= "col_1" | "col_2"
table ::= "table_1" | "table_2"
condition ::= column "=" number
number ::= "1" | "2"
"""


@pytest.fixture
def sample_sql_lark():
161
    return """
162
163
164
165
166
167
start: select_statement
select_statement: "SELECT" column "from" table "where" condition
column: "col_1" | "col_2"
table: "table_1" | "table_2"
condition: column "=" number
number: "1" | "2"
168
"""