mirror of
https://github.com/imartinez/privateGPT.git
synced 2026-07-17 01:48:03 +00:00
573 lines
18 KiB
Python
573 lines
18 KiB
Python
from typing import TYPE_CHECKING, Any
|
|
|
|
import pytest
|
|
from llama_index.core.base.llms.types import (
|
|
AudioBlock,
|
|
ImageBlock,
|
|
TextBlock,
|
|
)
|
|
from llama_index.core.llms import LLM, ChatMessage, MessageRole
|
|
from pydantic import Field
|
|
|
|
from private_gpt.components.chat.processors.chat_history.multimodality.audio_preprocessor import (
|
|
preprocess_audio_message,
|
|
)
|
|
from private_gpt.components.chat.processors.chat_history.multimodality.image_preprocessor import (
|
|
preprocess_image_message,
|
|
)
|
|
from private_gpt.components.chat.processors.chat_history.multimodality.multimodality_preprocessor import (
|
|
preprocess_multimodal_history,
|
|
preprocess_multimodal_message,
|
|
)
|
|
from private_gpt.components.chat.processors.chat_history.multimodality.utils import (
|
|
extract_audio_blocks,
|
|
extract_image_blocks,
|
|
requires_image_preprocessing,
|
|
)
|
|
from private_gpt.components.llm.custom.mock import FunctionCallingLLMMock
|
|
from private_gpt.components.multimodality.image_handler import (
|
|
ExtractionContent,
|
|
ExtractionEvaluation,
|
|
ExtractionStrategy,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from private_gpt.components.chat.processors.chat_history.multimodality.models import (
|
|
MultimodalProcessingResponse,
|
|
)
|
|
|
|
|
|
class MockLLM(FunctionCallingLLMMock):
|
|
responses: list[Any] = Field(default_factory=list)
|
|
call_count: int = Field(default=0)
|
|
messages_history: list[list[ChatMessage]] = Field(default_factory=list)
|
|
|
|
def __init__(self, responses: list[Any] | None = None, **kwargs: Any) -> None:
|
|
# Initialize parent class with only its expected parameters
|
|
super().__init__(**kwargs)
|
|
|
|
# Set our custom fields after parent initialization
|
|
if responses is not None:
|
|
self.responses = responses
|
|
self.responses = self.responses or []
|
|
self.call_count = 0
|
|
self.messages_history = []
|
|
|
|
async def astructured_chat(
|
|
self, output_cls: type, messages: list[ChatMessage], **kwargs: Any
|
|
) -> Any:
|
|
self.messages_history.append(messages)
|
|
if self.call_count < len(self.responses):
|
|
response = self.responses[self.call_count]
|
|
self.call_count += 1
|
|
return response
|
|
raise IndexError("No more mock responses available")
|
|
|
|
|
|
@pytest.fixture
|
|
def main_llm() -> LLM:
|
|
return MockLLM()
|
|
|
|
|
|
@pytest.fixture
|
|
def image_llm() -> LLM:
|
|
mock_responses = [
|
|
# Strategy inference
|
|
ExtractionStrategy(
|
|
type="table",
|
|
confidence=0.95,
|
|
language="en",
|
|
has_structure=True,
|
|
increase_contrast=False,
|
|
),
|
|
# Content extraction (complete)
|
|
ExtractionContent(
|
|
markdown="Describe these images",
|
|
is_complete=True,
|
|
),
|
|
# Evaluation (passes)
|
|
ExtractionEvaluation(
|
|
score=0.9,
|
|
issues_found=[],
|
|
),
|
|
]
|
|
|
|
llm = MockLLM(responses=mock_responses)
|
|
return llm
|
|
|
|
|
|
@pytest.fixture
|
|
def audio_llm() -> LLM:
|
|
from private_gpt.components.multimodality.audio_handler import (
|
|
TimestampSegment,
|
|
TranscriptionContent,
|
|
TranscriptionEvaluation,
|
|
TranscriptionStrategy,
|
|
)
|
|
|
|
mock_responses = [
|
|
TranscriptionStrategy(
|
|
type="speech",
|
|
confidence=0.9,
|
|
language="en",
|
|
has_multiple_speakers=False,
|
|
has_background_noise=False,
|
|
enhance_audio=False,
|
|
speaker_diarization=False,
|
|
),
|
|
TranscriptionContent(
|
|
timestamps=[
|
|
TimestampSegment(
|
|
start=0.0,
|
|
end=2.5,
|
|
speaker="Speaker 1",
|
|
text="Hello, how are you today?",
|
|
)
|
|
],
|
|
is_complete=True,
|
|
),
|
|
TranscriptionContent(
|
|
timestamps=[
|
|
TimestampSegment(
|
|
start=0.0,
|
|
end=2.5,
|
|
speaker="Speaker 1",
|
|
text="This is chunk two.",
|
|
)
|
|
],
|
|
is_complete=True,
|
|
),
|
|
TranscriptionContent(
|
|
timestamps=[
|
|
TimestampSegment(
|
|
start=0.0,
|
|
end=2.5,
|
|
speaker="Speaker 1",
|
|
text="This is chunk three.",
|
|
)
|
|
],
|
|
is_complete=True,
|
|
),
|
|
TranscriptionContent(
|
|
timestamps=[
|
|
TimestampSegment(
|
|
start=0.0,
|
|
end=2.5,
|
|
speaker="Speaker 1",
|
|
text="This is chunk four.",
|
|
)
|
|
],
|
|
is_complete=True,
|
|
),
|
|
TranscriptionContent(
|
|
timestamps=[
|
|
TimestampSegment(
|
|
start=0.0,
|
|
end=2.5,
|
|
speaker="Speaker 1",
|
|
text="This is chunk five.",
|
|
)
|
|
],
|
|
is_complete=True,
|
|
),
|
|
TranscriptionEvaluation(
|
|
score=0.9,
|
|
issues_found=[],
|
|
clarity=0.85,
|
|
),
|
|
]
|
|
|
|
llm = MockLLM(responses=mock_responses)
|
|
return llm
|
|
|
|
|
|
@pytest.fixture
|
|
def text_message() -> ChatMessage:
|
|
return ChatMessage(
|
|
role=MessageRole.USER,
|
|
blocks=[TextBlock(text="Hello, how are you?")],
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def image_message() -> ChatMessage:
|
|
return ChatMessage(
|
|
role=MessageRole.USER,
|
|
blocks=[
|
|
TextBlock(text="Describe these images"),
|
|
ImageBlock(url="https://picsum.photos/200/300"),
|
|
ImageBlock(url="https://picsum.photos/200/300"),
|
|
],
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def audio_message() -> ChatMessage:
|
|
return ChatMessage(
|
|
role=MessageRole.USER,
|
|
blocks=[
|
|
TextBlock(text="Transcribe this audio"),
|
|
AudioBlock(
|
|
url="https://commondatastorage.googleapis.com/codeskulptor-demos/DDR_assets/Kangaroo_MusiQue_-_The_Neverwritten_Role_Playing_Game.mp3"
|
|
),
|
|
],
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def multimodal_message() -> ChatMessage:
|
|
return ChatMessage(
|
|
role=MessageRole.USER,
|
|
blocks=[
|
|
TextBlock(text="Process this media"),
|
|
ImageBlock(url="https://picsum.photos/200/300"),
|
|
AudioBlock(
|
|
url="https://commondatastorage.googleapis.com/codeskulptor-demos/DDR_assets/Kangaroo_MusiQue_-_The_Neverwritten_Role_Playing_Game.mp3"
|
|
),
|
|
],
|
|
)
|
|
|
|
|
|
class TestImagePreprocessing:
|
|
@pytest.mark.asyncio
|
|
async def test_same_llm_no_preprocessing(
|
|
self, main_llm: LLM, image_message: ChatMessage
|
|
) -> None:
|
|
responses = []
|
|
async for response in preprocess_image_message(
|
|
main_llm, image_message, main_llm
|
|
):
|
|
responses.append(response)
|
|
|
|
result = responses[-1].message
|
|
assert result is image_message
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_different_llm_preprocessing_occurs(
|
|
self, main_llm: LLM, image_llm: LLM, image_message: ChatMessage
|
|
) -> None:
|
|
responses = []
|
|
async for response in preprocess_image_message(
|
|
main_llm, image_message, image_llm
|
|
):
|
|
responses.append(response)
|
|
|
|
result = responses[-1].message
|
|
assert result is not None
|
|
assert "images in their message" in str(result.blocks[-1].text)
|
|
assert result.role == MessageRole.USER
|
|
assert len(result.blocks) >= 2
|
|
assert isinstance(result.blocks[-1], TextBlock)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_text_message_passes_through(
|
|
self, main_llm: LLM, image_llm: LLM, text_message: ChatMessage
|
|
) -> None:
|
|
responses = []
|
|
async for response in preprocess_image_message(
|
|
main_llm, text_message, image_llm
|
|
):
|
|
responses.append(response)
|
|
|
|
result = responses[-1].message
|
|
assert result is text_message
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_missing_image_llm_raises_error(
|
|
self, main_llm: LLM, image_message: ChatMessage
|
|
) -> None:
|
|
with pytest.raises(
|
|
ValueError, match="Image blocks found but no image-capable LLM"
|
|
):
|
|
async for _ in preprocess_image_message(main_llm, image_message, None):
|
|
pass
|
|
|
|
|
|
class TestAudioPreprocessing:
|
|
@pytest.mark.asyncio
|
|
async def test_same_llm_no_preprocessing(
|
|
self, main_llm: LLM, audio_message: ChatMessage
|
|
) -> None:
|
|
responses = []
|
|
async for response in preprocess_audio_message(
|
|
main_llm, audio_message, main_llm
|
|
):
|
|
responses.append(response)
|
|
|
|
result = responses[-1].message
|
|
assert result is audio_message
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_different_llm_preprocessing_occurs(
|
|
self, main_llm: LLM, audio_llm: LLM, audio_message: ChatMessage
|
|
) -> None:
|
|
responses = []
|
|
async for response in preprocess_audio_message(
|
|
main_llm, audio_message, audio_llm
|
|
):
|
|
responses.append(response)
|
|
|
|
result = responses[-1].message
|
|
assert result is not None
|
|
assert "audios in their message" in str(result.blocks[-1].text)
|
|
assert result.role == MessageRole.USER
|
|
assert len(result.blocks) >= 2
|
|
assert isinstance(result.blocks[-1], TextBlock)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_text_message_passes_through(
|
|
self, main_llm: LLM, audio_llm: LLM, text_message: ChatMessage
|
|
) -> None:
|
|
responses = []
|
|
async for response in preprocess_audio_message(
|
|
main_llm, text_message, audio_llm
|
|
):
|
|
responses.append(response)
|
|
|
|
result = responses[-1].message
|
|
assert result is text_message
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_missing_audio_llm_raises_error(
|
|
self, main_llm: LLM, audio_message: ChatMessage
|
|
) -> None:
|
|
with pytest.raises(
|
|
ValueError, match="Audio blocks found but no audio-capable LLM"
|
|
):
|
|
async for _ in preprocess_audio_message(main_llm, audio_message, None):
|
|
pass
|
|
|
|
|
|
class TestMultimodalPreprocessing:
|
|
@pytest.mark.asyncio
|
|
async def test_same_llms_no_preprocessing(
|
|
self, main_llm: LLM, multimodal_message: ChatMessage
|
|
) -> None:
|
|
responses: list[MultimodalProcessingResponse] = []
|
|
async for response in preprocess_multimodal_message(
|
|
main_llm,
|
|
multimodal_message.model_copy(),
|
|
image_multimodal_llm=main_llm,
|
|
audio_multimodal_llm=main_llm,
|
|
):
|
|
responses.append(response)
|
|
|
|
result = responses[-1].modified_message
|
|
assert result.content == multimodal_message.content
|
|
assert result.role == multimodal_message.role
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_image_only_preprocessing(
|
|
self, main_llm: LLM, image_llm: LLM, image_message: ChatMessage
|
|
) -> None:
|
|
responses: list[MultimodalProcessingResponse] = []
|
|
async for response in preprocess_multimodal_message(
|
|
main_llm,
|
|
image_message,
|
|
image_multimodal_llm=image_llm,
|
|
audio_multimodal_llm=main_llm,
|
|
):
|
|
responses.append(response)
|
|
|
|
result = responses[-1].modified_message
|
|
assert len(result.blocks) >= 2
|
|
assert isinstance(result.blocks[-1], TextBlock)
|
|
assert "images in their message" in result.blocks[-1].text
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_audio_only_preprocessing(
|
|
self,
|
|
main_llm: LLM,
|
|
image_llm: LLM,
|
|
audio_llm: LLM,
|
|
multimodal_message: ChatMessage,
|
|
) -> None:
|
|
# Audio processing now yields failed status instead of raising directly
|
|
responses: list[MultimodalProcessingResponse] = []
|
|
async for response in preprocess_multimodal_message(
|
|
main_llm,
|
|
multimodal_message.model_copy(),
|
|
image_multimodal_llm=image_llm,
|
|
audio_multimodal_llm=audio_llm,
|
|
):
|
|
responses.append(response)
|
|
|
|
result = responses[-1].modified_message
|
|
assert len(result.blocks) >= 2
|
|
assert isinstance(result.blocks[-1], TextBlock)
|
|
assert "audios in their message" in result.blocks[-1].text
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_text_message_passes_through(
|
|
self,
|
|
main_llm: LLM,
|
|
image_llm: LLM,
|
|
audio_llm: LLM,
|
|
text_message: ChatMessage,
|
|
) -> None:
|
|
responses: list[MultimodalProcessingResponse] = []
|
|
async for response in preprocess_multimodal_message(
|
|
main_llm,
|
|
text_message,
|
|
image_multimodal_llm=image_llm,
|
|
audio_multimodal_llm=audio_llm,
|
|
):
|
|
responses.append(response)
|
|
|
|
result = responses[-1].modified_message
|
|
assert result.content == text_message.content
|
|
assert result.role == text_message.role
|
|
|
|
|
|
class TestHistoryPreprocessing:
|
|
@pytest.mark.asyncio
|
|
async def test_same_llm_no_preprocessing(
|
|
self, main_llm: LLM, text_message: ChatMessage, image_message: ChatMessage
|
|
) -> None:
|
|
history = [text_message, image_message]
|
|
|
|
responses: list[MultimodalProcessingResponse] = []
|
|
async for response in preprocess_multimodal_history(
|
|
main_llm,
|
|
history,
|
|
image_multimodal_llm=main_llm,
|
|
audio_multimodal_llm=main_llm,
|
|
):
|
|
responses.append(response)
|
|
|
|
result = responses[-1].chat_history
|
|
assert len(result) == len(history)
|
|
assert result[0].content == text_message.content
|
|
assert result[1].content == image_message.content
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_processes_most_recent_user_message(
|
|
self,
|
|
main_llm: LLM,
|
|
image_llm: LLM,
|
|
text_message: ChatMessage,
|
|
image_message: ChatMessage,
|
|
) -> None:
|
|
history = [
|
|
ChatMessage(
|
|
role=MessageRole.USER,
|
|
content="First",
|
|
blocks=[ImageBlock(url="https://picsum.photos/300/300")],
|
|
),
|
|
ChatMessage(role=MessageRole.ASSISTANT, content="Response", blocks=[]),
|
|
image_message,
|
|
]
|
|
|
|
responses: list[MultimodalProcessingResponse] = []
|
|
async for response in preprocess_multimodal_history(
|
|
main_llm, history, image_multimodal_llm=image_llm, audio_multimodal_llm=None
|
|
):
|
|
responses.append(response)
|
|
|
|
result = responses[-1].chat_history
|
|
assert len(result) == 3
|
|
assert len(result[2].blocks) >= 2
|
|
assert isinstance(result[2].blocks[-1], TextBlock)
|
|
assert "images in their message" in result[2].blocks[-1].text
|
|
|
|
assert result[0].content == "First"
|
|
assert all(isinstance(block, TextBlock) for block in result[0].blocks)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_empty_history_returns_none(
|
|
self, main_llm: LLM, image_llm: LLM
|
|
) -> None:
|
|
responses: list[MultimodalProcessingResponse] = []
|
|
async for response in preprocess_multimodal_history(
|
|
main_llm, None, image_multimodal_llm=image_llm, audio_multimodal_llm=None
|
|
):
|
|
responses.append(response)
|
|
|
|
assert len(responses) == 1
|
|
assert responses[0].chat_history is None
|
|
|
|
responses: list[MultimodalProcessingResponse] = []
|
|
async for response in preprocess_multimodal_history(
|
|
main_llm, [], image_multimodal_llm=image_llm, audio_multimodal_llm=None
|
|
):
|
|
responses.append(response)
|
|
|
|
assert len(responses) == 1
|
|
assert responses[0].chat_history == []
|
|
|
|
|
|
class TestTypeChecking:
|
|
def test_llm_objects_identity_comparison(
|
|
self, main_llm: LLM, image_llm: LLM
|
|
) -> None:
|
|
assert not requires_image_preprocessing(main_llm, main_llm)
|
|
assert requires_image_preprocessing(main_llm, image_llm)
|
|
assert requires_image_preprocessing(main_llm, None)
|
|
|
|
def test_extract_functions_return_correct_types(
|
|
self, multimodal_message: ChatMessage
|
|
) -> None:
|
|
images = extract_image_blocks(multimodal_message)
|
|
audios = extract_audio_blocks(multimodal_message)
|
|
|
|
assert isinstance(images, list)
|
|
assert isinstance(audios, list)
|
|
assert len(images) == 1
|
|
assert len(audios) == 1
|
|
assert isinstance(images[0], ImageBlock)
|
|
assert isinstance(audios[0], AudioBlock)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("same_image_llm", "same_audio_llm", "expected_preprocessing"),
|
|
[
|
|
(True, True, False),
|
|
(True, False, True),
|
|
(False, True, True),
|
|
(False, False, True),
|
|
],
|
|
)
|
|
class TestParametrizedPreprocessing:
|
|
@pytest.mark.asyncio
|
|
async def test_preprocessing_decision_logic(
|
|
self,
|
|
main_llm: LLM,
|
|
image_llm: LLM,
|
|
audio_llm: LLM,
|
|
multimodal_message: ChatMessage,
|
|
same_image_llm: bool,
|
|
same_audio_llm: bool,
|
|
expected_preprocessing: bool,
|
|
) -> None:
|
|
actual_image_llm = main_llm if same_image_llm else image_llm
|
|
actual_audio_llm = main_llm if same_audio_llm else audio_llm
|
|
|
|
responses: list[MultimodalProcessingResponse] = []
|
|
async for response in preprocess_multimodal_message(
|
|
main_llm,
|
|
multimodal_message.model_copy(),
|
|
image_multimodal_llm=actual_image_llm,
|
|
audio_multimodal_llm=actual_audio_llm,
|
|
):
|
|
responses.append(response)
|
|
|
|
result = responses[-1].modified_message
|
|
if expected_preprocessing:
|
|
assert result.content != multimodal_message.content
|
|
if not same_image_llm:
|
|
assert len(result.blocks) >= 2
|
|
assert any(
|
|
"images in their message" in block.text
|
|
for block in result.blocks
|
|
if isinstance(block, TextBlock)
|
|
)
|
|
if not same_audio_llm:
|
|
assert any(
|
|
"audios in their message" in block.text
|
|
for block in result.blocks
|
|
if isinstance(block, TextBlock)
|
|
)
|
|
else:
|
|
assert result.content == multimodal_message.content
|
|
assert result.role == multimodal_message.role
|