image to text support, closes #134

3 files changed, 257 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
diff --git a/tests/resources/binary_file.bin b/tests/resources/binary_file.bin
new file mode 100644
index 0000000..acd83ea
--- /dev/null
+++ b/tests/resources/binary_file.bin
@@ -0,0 +1,2 @@
+�#n@�E[+)�W~���Eo{��ZL��>�]��^�����ʒ��Z�xv�"�E%,�b�oqq�ȸsv�$)b��y�7¶|��SC��.�	!�UL����3��S��b�u�$cj����ӧ����	�}��D�>�|�����E
+rc(�*(�����hW�GZ��?��쮘po�$�B\��vz풮ƌ:'�=$o6�v�~p$o5h�a��M�����Sy(9��ֺf�`O�a��|�>kL���
\ No newline at end of file
diff --git a/tests/resources/image_file.jpg b/tests/resources/image_file.jpg
new file mode 100644
index 0000000..3433468
--- /dev/null
+++ b/tests/resources/image_file.jpg
@@ -0,0 +1 @@
+image data