blob: 9acfecbd8d8e3ef7e2fc1a01fb89ef8bbd2811a3 (
plain)
1
2
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
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
|
from utils import parse_chat_messages
import os
curr_dir = os.path.dirname(__file__)
def strip_text(txt):
txt = txt.strip()
lines = txt.splitlines()
return "\n".join([line.lstrip() for line in lines])
def test_parse_user_message():
chat_content = strip_text(
"""
>>> user
generate lorem ipsum
""")
messages = parse_chat_messages(chat_content)
assert 1 == len(messages)
assert 'user' == messages[0]['role']
assert 'generate lorem ipsum' == messages[0]['content']
def test_parse_system_message():
chat_content = strip_text("""
>>> system
you are general assystant
>>> user
generate lorem ipsum
""")
messages = parse_chat_messages(chat_content)
assert 2 == len(messages)
assert 'system' == messages[0]['role']
assert 'you are general assystant' == messages[0]['content']
assert 'user' == messages[1]['role']
assert 'generate lorem ipsum' == messages[1]['content']
def test_parse_assistant_message():
chat_content = strip_text("""
>>> user
generate lorem ipsum
<<< assistant
bla bla bla
>>> user
again
""")
messages = parse_chat_messages(chat_content)
assert 3 == len(messages)
assert 'user' == messages[0]['role']
assert 'generate lorem ipsum' == messages[0]['content']
assert 'assistant' == messages[1]['role']
assert 'bla bla bla' == messages[1]['content']
assert 'user' == messages[2]['role']
assert 'again' == messages[2]['content']
def test_parse_include_single_file_message():
chat_content = strip_text(f"""
>>> user
translate to human language
>>> include
{curr_dir}/resources/test1.include.txt
""")
messages = parse_chat_messages(chat_content)
assert 2 == len(messages)
assert 'user' == messages[0]['role']
assert 'translate to human language' == messages[0]['content']
assert 'user' == messages[1]['role']
expected_content = strip_text(f"""
==> {curr_dir}/resources/test1.include.txt <==
hello world
""")
assert expected_content == messages[1]['content']
def test_parse_include_multiple_files_message():
chat_content = strip_text(f"""
>>> user
translate to human language
>>> include
{curr_dir}/resources/test1.include.txt
{curr_dir}/resources/test2.include.txt
""")
messages = parse_chat_messages(chat_content)
assert 2 == len(messages)
assert 'user' == messages[0]['role']
assert 'translate to human language' == messages[0]['content']
assert 'user' == messages[1]['role']
expected_content = strip_text(f"""
==> {curr_dir}/resources/test1.include.txt <==
hello world
==> {curr_dir}/resources/test2.include.txt <==
vim is awesome
""")
assert expected_content == messages[1]['content']
def test_parse_include_glob_files_message():
chat_content = strip_text(f"""
>>> user
translate to human language
>>> include
{curr_dir}/**/*.include.txt
""")
messages = parse_chat_messages(chat_content)
assert 2 == len(messages)
assert 'user' == messages[0]['role']
assert 'translate to human language' == messages[0]['content']
assert 'user' == messages[1]['role']
expected_content = strip_text(f"""
==> {curr_dir}/resources/test1.include.txt <==
hello world
==> {curr_dir}/resources/test2.include.txt <==
vim is awesome
""")
assert expected_content == messages[1]['content']
def test_parse_include_image_message():
# TODO
pass
def test_parse_include_image_with_files_message():
# TODO
pass
|
