fix: prevent crash in parse_vision_messages when vision is disabled (#5487)
This commit is contained in:
+13
-4
@@ -179,11 +179,20 @@ def parse_vision_messages(messages, llm=None, vision_details="auto"):
|
|||||||
|
|
||||||
# Handle message content
|
# Handle message content
|
||||||
if isinstance(msg["content"], list):
|
if isinstance(msg["content"], list):
|
||||||
# Multiple image URLs in content
|
if llm is None:
|
||||||
description = get_image_description(msg, llm, vision_details)
|
text_parts = [
|
||||||
returned_messages.append({"role": msg["role"], "content": description})
|
part["text"] for part in msg["content"]
|
||||||
|
if isinstance(part, dict) and part.get("type") == "text"
|
||||||
|
]
|
||||||
|
if not text_parts:
|
||||||
|
continue
|
||||||
|
returned_messages.append({"role": msg["role"], "content": " ".join(text_parts)})
|
||||||
|
else:
|
||||||
|
description = get_image_description(msg, llm, vision_details)
|
||||||
|
returned_messages.append({"role": msg["role"], "content": description})
|
||||||
elif isinstance(msg["content"], dict) and msg["content"].get("type") == "image_url":
|
elif isinstance(msg["content"], dict) and msg["content"].get("type") == "image_url":
|
||||||
# Single image content
|
if llm is None:
|
||||||
|
continue
|
||||||
image_url = msg["content"]["image_url"]["url"]
|
image_url = msg["content"]["image_url"]["url"]
|
||||||
try:
|
try:
|
||||||
description = get_image_description(image_url, llm, vision_details)
|
description = get_image_description(image_url, llm, vision_details)
|
||||||
|
|||||||
@@ -1,5 +1,60 @@
|
|||||||
import pytest
|
import pytest
|
||||||
from mem0.memory.utils import remove_spaces_from_entities, sanitize_relationship_for_cypher
|
from unittest.mock import Mock
|
||||||
|
|
||||||
|
from mem0.memory.utils import parse_vision_messages, remove_spaces_from_entities, sanitize_relationship_for_cypher
|
||||||
|
|
||||||
|
|
||||||
|
class TestParseVisionMessages:
|
||||||
|
def test_multimodal_list_without_llm_extracts_text(self):
|
||||||
|
messages = [
|
||||||
|
{"role": "user", "content": [
|
||||||
|
{"type": "text", "text": "What is this?"},
|
||||||
|
{"type": "image_url", "image_url": {"url": "https://example.com/img.png"}},
|
||||||
|
]},
|
||||||
|
]
|
||||||
|
result = parse_vision_messages(messages, llm=None)
|
||||||
|
assert len(result) == 1
|
||||||
|
assert result[0]["role"] == "user"
|
||||||
|
assert result[0]["content"] == "What is this?"
|
||||||
|
|
||||||
|
def test_image_dict_without_llm_is_skipped(self):
|
||||||
|
messages = [
|
||||||
|
{"role": "user", "content": {"type": "image_url", "image_url": {"url": "https://example.com/img.png"}}},
|
||||||
|
{"role": "user", "content": "hello"},
|
||||||
|
]
|
||||||
|
result = parse_vision_messages(messages, llm=None)
|
||||||
|
assert len(result) == 1
|
||||||
|
assert result[0]["content"] == "hello"
|
||||||
|
|
||||||
|
def test_multimodal_with_llm_calls_generate_response(self):
|
||||||
|
mock_llm = Mock()
|
||||||
|
mock_llm.generate_response.return_value = "A photo of a cat"
|
||||||
|
messages = [
|
||||||
|
{"role": "user", "content": [
|
||||||
|
{"type": "text", "text": "Describe this"},
|
||||||
|
{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}},
|
||||||
|
]},
|
||||||
|
]
|
||||||
|
result = parse_vision_messages(messages, llm=mock_llm, vision_details="auto")
|
||||||
|
assert result[0]["content"] == "A photo of a cat"
|
||||||
|
mock_llm.generate_response.assert_called_once()
|
||||||
|
|
||||||
|
def test_image_only_list_without_llm_is_skipped(self):
|
||||||
|
messages = [
|
||||||
|
{"role": "user", "content": [
|
||||||
|
{"type": "image_url", "image_url": {"url": "https://example.com/img.png"}},
|
||||||
|
]},
|
||||||
|
]
|
||||||
|
result = parse_vision_messages(messages, llm=None)
|
||||||
|
assert result == []
|
||||||
|
|
||||||
|
def test_plain_text_messages_pass_through(self):
|
||||||
|
messages = [
|
||||||
|
{"role": "system", "content": "You are helpful."},
|
||||||
|
{"role": "user", "content": "Hello"},
|
||||||
|
]
|
||||||
|
result = parse_vision_messages(messages, llm=None)
|
||||||
|
assert result == messages
|
||||||
|
|
||||||
|
|
||||||
class TestRemoveSpacesFromEntities:
|
class TestRemoveSpacesFromEntities:
|
||||||
|
|||||||
Reference in New Issue
Block a user