summaryrefslogtreecommitdiff
path: root/tests/chat_test.py
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--tests/chat_test.py310
1 files changed, 254 insertions, 56 deletions
diff --git a/tests/chat_test.py b/tests/chat_test.py
index 9acfecb..e29368c 100644
--- a/tests/chat_test.py
+++ b/tests/chat_test.py
@@ -15,10 +15,18 @@ def test_parse_user_message():
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']
+ actual_messages = parse_chat_messages(chat_content)
+ assert [
+ {
+ 'role': 'user',
+ 'content': [
+ {
+ 'type': 'text',
+ 'text': 'generate lorem ipsum',
+ },
+ ],
+ },
+ ] == actual_messages
def test_parse_system_message():
@@ -31,12 +39,56 @@ def test_parse_system_message():
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']
+ actual_messages = parse_chat_messages(chat_content)
+ assert [
+ {
+ 'role': 'system',
+ 'content': [
+ {
+ 'type': 'text',
+ 'text': 'you are general assystant',
+ },
+ ],
+ },
+ {
+ 'role': 'user',
+ 'content': [
+ {
+ 'type': 'text',
+ 'text': 'generate lorem ipsum',
+ },
+ ],
+ },
+ ] == actual_messages
+
+
+def test_parse_two_user_messages():
+ chat_content = strip_text(
+ """
+ >>> user
+
+ generate lorem ipsum
+
+ >>> user
+
+ in english
+ """)
+ actual_messages = parse_chat_messages(chat_content)
+ assert [
+ {
+ 'role': 'user',
+ 'content': [
+ {
+ 'type': 'text',
+ 'text': 'generate lorem ipsum',
+ },
+ {
+ 'type': 'text',
+ 'text': 'in english',
+ },
+ ],
+ },
+ ] == actual_messages
def test_parse_assistant_message():
chat_content = strip_text("""
@@ -52,14 +104,36 @@ def test_parse_assistant_message():
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']
+ actual_messages = parse_chat_messages(chat_content)
+ assert [
+ {
+ 'role': 'user',
+ 'content': [
+ {
+ 'type': 'text',
+ 'text': 'generate lorem ipsum',
+ },
+ ],
+ },
+ {
+ 'role': 'assistant',
+ 'content': [
+ {
+ 'type': 'text',
+ 'text': 'bla bla bla',
+ },
+ ],
+ },
+ {
+ 'role': 'user',
+ 'content': [
+ {
+ 'type': 'text',
+ 'text': 'again',
+ },
+ ],
+ },
+ ] == actual_messages
def test_parse_include_single_file_message():
chat_content = strip_text(f"""
@@ -70,17 +144,50 @@ def test_parse_include_single_file_message():
>>> include
{curr_dir}/resources/test1.include.txt
+
+ <<< assistant
+
+ it already is in human language
+
+ >>> user
+
+ try harder
""")
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']
+ actual_messages = parse_chat_messages(chat_content)
+ assert [
+ {
+ 'role': 'user',
+ 'content': [
+ {
+ 'type': 'text',
+ 'text': 'translate to human language',
+ },
+ {
+ 'type': 'text',
+ 'text': f'==> {curr_dir}/resources/test1.include.txt <==\nhello world',
+ },
+ ],
+ },
+ {
+ 'role': 'assistant',
+ 'content': [
+ {
+ 'type': 'text',
+ 'text': 'it already is in human language',
+ },
+ ],
+ },
+ {
+ 'role': 'user',
+ 'content': [
+ {
+ 'type': 'text',
+ 'text': 'try harder',
+ },
+ ],
+ },
+ ] == actual_messages
def test_parse_include_multiple_files_message():
chat_content = strip_text(f"""
@@ -94,18 +201,26 @@ def test_parse_include_multiple_files_message():
{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']
+ actual_messages = parse_chat_messages(chat_content)
+ assert [
+ {
+ 'role': 'user',
+ 'content': [
+ {
+ 'type': 'text',
+ 'text': 'translate to human language',
+ },
+ {
+ 'type': 'text',
+ 'text': f'==> {curr_dir}/resources/test1.include.txt <==\nhello world',
+ },
+ {
+ 'type': 'text',
+ 'text': f'==> {curr_dir}/resources/test2.include.txt <==\nvim is awesome',
+ },
+ ],
+ },
+ ] == actual_messages
def test_parse_include_glob_files_message():
chat_content = strip_text(f"""
@@ -117,24 +232,107 @@ def test_parse_include_glob_files_message():
{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']
+ actual_messages = parse_chat_messages(chat_content)
+ assert [
+ {
+ 'role': 'user',
+ 'content': [
+ {
+ 'type': 'text',
+ 'text': 'translate to human language',
+ },
+ {
+ 'type': 'text',
+ 'text': f'==> {curr_dir}/resources/test1.include.txt <==\nhello world',
+ },
+ {
+ 'type': 'text',
+ 'text': f'==> {curr_dir}/resources/test2.include.txt <==\nvim is awesome',
+ },
+ ],
+ },
+ ] == actual_messages
def test_parse_include_image_message():
- # TODO
- pass
+ chat_content = strip_text(f"""
+ >>> user
+
+ what is on the image?
+
+ >>> include
+
+ {curr_dir}/**/*.jpg
+ """)
+ actual_messages = parse_chat_messages(chat_content)
+ assert [
+ {
+ 'role': 'user',
+ 'content': [
+ {
+ 'type': 'text',
+ 'text': 'what is on the image?',
+ },
+ {
+ 'type': 'image_url',
+ 'image_url': {
+ 'url': 'data:image/jpg;base64,aW1hZ2UgZGF0YQo='
+ },
+ },
+ ],
+ },
+ ] == actual_messages
def test_parse_include_image_with_files_message():
- # TODO
- pass
+ chat_content = strip_text(f"""
+ >>> include
+
+ {curr_dir}/resources/test1.include.txt
+ {curr_dir}/resources/image_file.jpg
+ {curr_dir}/resources/test2.include.txt
+ """)
+ actual_messages = parse_chat_messages(chat_content)
+ assert [
+ {
+ 'role': 'user',
+ 'content': [
+ {
+ 'type': 'text',
+ 'text': f'==> {curr_dir}/resources/test1.include.txt <==\nhello world',
+ },
+ {
+ 'type': 'image_url',
+ 'image_url': {
+ 'url': 'data:image/jpg;base64,aW1hZ2UgZGF0YQo='
+ },
+ },
+ {
+ 'type': 'text',
+ 'text': f'==> {curr_dir}/resources/test2.include.txt <==\nvim is awesome',
+ },
+ ],
+ },
+ ] == actual_messages
+
+def test_parse_include_unsupported_binary_file():
+ chat_content = strip_text(f"""
+ >>> include
+
+ {curr_dir}/resources/binary_file.bin
+ {curr_dir}/resources/test1.include.txt
+ """)
+ actual_messages = parse_chat_messages(chat_content)
+ assert [
+ {
+ 'role': 'user',
+ 'content': [
+ {
+ 'type': 'text',
+ 'text': f'==> {curr_dir}/resources/binary_file.bin <==\nBinary file, cannot display',
+ },
+ {
+ 'type': 'text',
+ 'text': f'==> {curr_dir}/resources/test1.include.txt <==\nhello world',
+ },
+ ],
+ },
+ ] == actual_messages
generated by cgit v1.2.3 (git 2.25.1) at 2025-12-10 11:42:05 +0000