attach images to message

This commit is contained in:
garylin2099 2024-08-14 20:12:17 +08:00
parent dce5502c07
commit e9984f2bf8
6 changed files with 83 additions and 33 deletions

View file

@ -10,8 +10,9 @@ import pytest
from metagpt.configs.compress_msg_config import CompressType
from metagpt.configs.llm_config import LLMConfig
from metagpt.const import IMAGES
from metagpt.provider.base_llm import BaseLLM
from metagpt.schema import Message
from metagpt.schema import AIMessage, Message, UserMessage
from tests.metagpt.provider.mock_llm_config import mock_llm_config
from tests.metagpt.provider.req_resp_const import (
default_resp_cont,
@ -163,3 +164,41 @@ def test_compress_messages_long_no_sys_msg(compress_type):
print(compressed)
assert compressed
assert len(compressed[0]["content"]) < len(messages[0]["content"])
def test_format_msg(mocker):
base_llm = MockBaseLLM()
messages = [UserMessage(content="req"), AIMessage(content="rsp")]
formatted_msgs = base_llm.format_msg(messages)
assert formatted_msgs == [{"role": "user", "content": "req"}, {"role": "assistant", "content": "rsp"}]
def test_format_msg_w_images(mocker):
base_llm = MockBaseLLM()
base_llm.config.model = "gpt-4o"
msg_w_images = UserMessage(content="req1")
msg_w_images.add_metadata(IMAGES, ["base64 string 1", "base64 string 2"])
msg_w_empty_images = UserMessage(content="req2")
msg_w_empty_images.add_metadata(IMAGES, [])
messages = [
msg_w_images, # should be converted
AIMessage(content="rsp"),
msg_w_empty_images, # should not be converted
]
formatted_msgs = base_llm.format_msg(messages)
assert formatted_msgs == [
{
"role": "user",
"content": [
{"type": "text", "text": "req1"},
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,base64 string 1"}},
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,base64 string 2"}},
],
},
{"role": "assistant", "content": "rsp"},
{"role": "user", "content": "req2"},
]
if name == "__main__":
pytest.main([__file__, "-s"])

View file

@ -29,9 +29,9 @@ from metagpt.utils.common import (
awrite,
check_cmd_exists,
concat_namespace,
extract_and_encode_images,
extract_image_paths,
import_class_inst,
is_support_image_input,
parse_recipient,
print_members,
read_file_block,
@ -231,9 +231,8 @@ def test_extract_image_paths():
assert not extract_image_paths(content)
def test_is_support_image_input():
assert is_support_image_input("gpt-4o-2024-08-06")
assert not is_support_image_input("deepseek-coder")
def test_extract_and_encode_images():
assert not extract_and_encode_images("a non-existing.jpg")
if __name__ == "__main__":