519 lines
19 KiB
Python
519 lines
19 KiB
Python
# Copyright (c) "Neo4j"
|
|
# Neo4j Sweden AB [https://neo4j.com]
|
|
# #
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
# #
|
|
# https://www.apache.org/licenses/LICENSE-2.0
|
|
# #
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
import warnings
|
|
from typing import Any, List, Optional, cast
|
|
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
|
|
|
import httpx
|
|
import pytest
|
|
from neo4j_graphrag.exceptions import LLMGenerationError
|
|
from neo4j_graphrag.llm import LLMResponse, MistralAILLM
|
|
from neo4j_graphrag.types import LLMMessage
|
|
from neo4j_graphrag.utils.rate_limit import NoOpRateLimitHandler
|
|
from pydantic import BaseModel, ConfigDict
|
|
|
|
|
|
# Mock SDKError for testing
|
|
class MockSDKError(Exception):
|
|
"""Mock SDKError for testing purposes."""
|
|
|
|
def __init__(
|
|
self, message: str, raw_response: Optional[httpx.Response] = None
|
|
) -> None:
|
|
super().__init__(message)
|
|
self.raw_response = raw_response
|
|
|
|
|
|
def _as_mock(value: Any) -> MagicMock:
|
|
return cast(MagicMock, value)
|
|
|
|
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral", None)
|
|
def test_mistralai_llm_missing_dependency() -> None:
|
|
with pytest.raises(ImportError):
|
|
MistralAILLM(model_name="mistral-model")
|
|
|
|
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
def test_mistralai_llm_invoke(mock_mistral: Mock) -> None:
|
|
mock_mistral_instance = mock_mistral.return_value
|
|
|
|
chat_response_mock = MagicMock()
|
|
chat_response_mock.choices = [
|
|
MagicMock(message=MagicMock(content="mistral response"))
|
|
]
|
|
|
|
mock_mistral_instance.chat.complete.return_value = chat_response_mock
|
|
|
|
llm = MistralAILLM(model_name="mistral-model")
|
|
|
|
res = llm.invoke("some input")
|
|
|
|
assert isinstance(res, LLMResponse)
|
|
assert res.content == "mistral response"
|
|
|
|
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
def test_mistralai_llm_invoke_with_message_history(mock_mistral: Mock) -> None:
|
|
mock_mistral_instance = mock_mistral.return_value
|
|
chat_response_mock = MagicMock()
|
|
chat_response_mock.choices = [
|
|
MagicMock(message=MagicMock(content="mistral response"))
|
|
]
|
|
mock_mistral_instance.chat.complete.return_value = chat_response_mock
|
|
model = "mistral-model"
|
|
system_instruction = "You are a helpful assistant."
|
|
|
|
llm = MistralAILLM(model_name=model)
|
|
|
|
message_history: List[LLMMessage] = [
|
|
{"role": "user", "content": "When does the sun come up in the summer?"},
|
|
{"role": "assistant", "content": "Usually around 6am."},
|
|
]
|
|
question = "What about next season?"
|
|
res = llm.invoke(question, message_history, system_instruction=system_instruction)
|
|
|
|
assert isinstance(res, LLMResponse)
|
|
assert res.content == "mistral response"
|
|
messages: List[LLMMessage] = [{"role": "system", "content": system_instruction}]
|
|
messages.extend(message_history)
|
|
messages.append({"role": "user", "content": question})
|
|
_as_mock(llm.client.chat.complete).assert_called_once_with(
|
|
messages=messages,
|
|
model=model,
|
|
)
|
|
|
|
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
def test_mistralai_llm_invoke_with_message_history_and_system_instruction(
|
|
mock_mistral: Mock,
|
|
) -> None:
|
|
mock_mistral_instance = mock_mistral.return_value
|
|
chat_response_mock = MagicMock()
|
|
chat_response_mock.choices = [
|
|
MagicMock(message=MagicMock(content="mistral response"))
|
|
]
|
|
mock_mistral_instance.chat.complete.return_value = chat_response_mock
|
|
model = "mistral-model"
|
|
system_instruction = "You are a helpful assistant."
|
|
llm = MistralAILLM(model_name=model)
|
|
message_history: List[LLMMessage] = [
|
|
{"role": "user", "content": "When does the sun come up in the summer?"},
|
|
{"role": "assistant", "content": "Usually around 6am."},
|
|
]
|
|
question = "What about next season?"
|
|
|
|
# first invocation - initial instructions
|
|
res = llm.invoke(question, message_history, system_instruction=system_instruction)
|
|
assert isinstance(res, LLMResponse)
|
|
assert res.content == "mistral response"
|
|
messages: List[LLMMessage] = [{"role": "system", "content": system_instruction}]
|
|
messages.extend(message_history)
|
|
messages.append({"role": "user", "content": question})
|
|
_as_mock(llm.client.chat.complete).assert_called_once_with(
|
|
messages=messages,
|
|
model=model,
|
|
)
|
|
|
|
assert _as_mock(llm.client.chat.complete).call_count == 1
|
|
|
|
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
def test_mistralai_llm_invoke_with_message_history_validation_error(
|
|
mock_mistral: Mock,
|
|
) -> None:
|
|
mock_mistral_instance = mock_mistral.return_value
|
|
chat_response_mock = MagicMock()
|
|
chat_response_mock.choices = [
|
|
MagicMock(message=MagicMock(content="mistral response"))
|
|
]
|
|
mock_mistral_instance.chat.complete.return_value = chat_response_mock
|
|
model = "mistral-model"
|
|
system_instruction = "You are a helpful assistant."
|
|
|
|
llm = MistralAILLM(model_name=model, system_instruction=system_instruction)
|
|
|
|
message_history = [
|
|
{"role": "user", "content": "When does the sun come up in the summer?"},
|
|
{"role": "monkey", "content": "Usually around 6am."},
|
|
]
|
|
question = "What about next season?"
|
|
|
|
with pytest.raises(LLMGenerationError) as exc_info:
|
|
llm.invoke(question, message_history) # type: ignore
|
|
assert "Input should be 'user', 'assistant' or 'system" in str(exc_info.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
async def test_mistralai_llm_ainvoke(mock_mistral: Mock) -> None:
|
|
mock_mistral_instance = mock_mistral.return_value
|
|
|
|
async def mock_complete_async(*_args: Any, **_kwargs: Any) -> MagicMock:
|
|
chat_response_mock = MagicMock()
|
|
chat_response_mock.choices = [
|
|
MagicMock(message=MagicMock(content="async mistral response"))
|
|
]
|
|
return chat_response_mock
|
|
|
|
mock_mistral_instance.chat.complete_async = mock_complete_async
|
|
|
|
llm = MistralAILLM(model_name="mistral-model")
|
|
|
|
res = await llm.ainvoke("some input")
|
|
|
|
assert isinstance(res, LLMResponse)
|
|
assert res.content == "async mistral response"
|
|
|
|
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.SDKError", MockSDKError)
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
def test_mistralai_llm_invoke_sdkerror(mock_mistral: Mock) -> None:
|
|
mock_mistral_instance = mock_mistral.return_value
|
|
raw_response = httpx.Response(status_code=500)
|
|
mock_mistral_instance.chat.complete.side_effect = MockSDKError(
|
|
"Some error", raw_response=raw_response
|
|
)
|
|
|
|
llm = MistralAILLM(model_name="mistral-model")
|
|
|
|
with pytest.raises(LLMGenerationError):
|
|
llm.invoke("some input")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.SDKError", MockSDKError)
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
async def test_mistralai_llm_ainvoke_sdkerror(mock_mistral: Mock) -> None:
|
|
mock_mistral_instance = mock_mistral.return_value
|
|
|
|
async def mock_complete_async(*args: Any, **kwargs: Any) -> None:
|
|
raw_response = httpx.Response(status_code=500)
|
|
raise MockSDKError("Some async error", raw_response=raw_response)
|
|
|
|
mock_mistral_instance.chat.complete_async = mock_complete_async
|
|
|
|
llm = MistralAILLM(model_name="mistral-model")
|
|
|
|
with pytest.raises(LLMGenerationError):
|
|
await llm.ainvoke("some input")
|
|
|
|
|
|
# V2 Interface Tests (List[LLMMessage] input)
|
|
|
|
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
def test_mistralai_llm_invoke_v2_happy_path(mock_mistral: Mock) -> None:
|
|
"""Test V2 interface invoke method with List[LLMMessage] input."""
|
|
mock_mistral_instance = mock_mistral.return_value
|
|
chat_response_mock = MagicMock()
|
|
chat_response_mock.choices = [
|
|
MagicMock(message=MagicMock(content="mistral v2 response"))
|
|
]
|
|
mock_mistral_instance.chat.complete.return_value = chat_response_mock
|
|
|
|
messages: List[LLMMessage] = [
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
{"role": "user", "content": "What is machine learning?"},
|
|
]
|
|
|
|
llm = MistralAILLM(model_name="mistral-model")
|
|
response = llm.invoke(messages)
|
|
|
|
assert isinstance(response, LLMResponse)
|
|
assert response.content == "mistral v2 response"
|
|
|
|
# Verify the correct method was called
|
|
_as_mock(llm.client.chat.complete).assert_called_once()
|
|
call_args = _as_mock(llm.client.chat.complete).call_args[1]
|
|
assert call_args["model"] == "mistral-model"
|
|
assert len(call_args["messages"]) == 2
|
|
|
|
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
def test_mistralai_llm_invoke_v2_with_conversation_history(mock_mistral: Mock) -> None:
|
|
"""Test V2 interface invoke method with complex conversation history."""
|
|
mock_mistral_instance = mock_mistral.return_value
|
|
chat_response_mock = MagicMock()
|
|
chat_response_mock.choices = [
|
|
MagicMock(message=MagicMock(content="mistral conversation response"))
|
|
]
|
|
mock_mistral_instance.chat.complete.return_value = chat_response_mock
|
|
|
|
messages: List[LLMMessage] = [
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
{"role": "user", "content": "Tell me about Python."},
|
|
{"role": "assistant", "content": "Python is a programming language."},
|
|
{"role": "user", "content": "What about its history?"},
|
|
]
|
|
|
|
llm = MistralAILLM(model_name="mistral-model")
|
|
response = llm.invoke(messages)
|
|
|
|
assert isinstance(response, LLMResponse)
|
|
assert response.content == "mistral conversation response"
|
|
|
|
# Verify the correct number of messages were passed
|
|
_as_mock(llm.client.chat.complete).assert_called_once()
|
|
call_args = _as_mock(llm.client.chat.complete).call_args[1]
|
|
assert len(call_args["messages"]) == 4
|
|
|
|
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
def test_mistralai_llm_invoke_v2_no_system_message(mock_mistral: Mock) -> None:
|
|
"""Test V2 interface invoke method without system message."""
|
|
mock_mistral_instance = mock_mistral.return_value
|
|
chat_response_mock = MagicMock()
|
|
chat_response_mock.choices = [
|
|
MagicMock(message=MagicMock(content="mistral no system response"))
|
|
]
|
|
mock_mistral_instance.chat.complete.return_value = chat_response_mock
|
|
|
|
messages: List[LLMMessage] = [
|
|
{"role": "user", "content": "What is the capital of France?"},
|
|
]
|
|
|
|
llm = MistralAILLM(model_name="mistral-model")
|
|
response = llm.invoke(messages)
|
|
|
|
assert isinstance(response, LLMResponse)
|
|
assert response.content == "mistral no system response"
|
|
|
|
# Verify only user message was passed
|
|
_as_mock(llm.client.chat.complete).assert_called_once()
|
|
call_args = _as_mock(llm.client.chat.complete).call_args[1]
|
|
assert len(call_args["messages"]) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
async def test_mistralai_llm_ainvoke_v2_happy_path(mock_mistral: Mock) -> None:
|
|
"""Test V2 interface async invoke method with List[LLMMessage] input."""
|
|
mock_mistral_instance = mock_mistral.return_value
|
|
|
|
async def mock_complete_async(*_args: Any, **_kwargs: Any) -> MagicMock:
|
|
chat_response_mock = MagicMock()
|
|
chat_response_mock.choices = [
|
|
MagicMock(message=MagicMock(content="async mistral v2 response"))
|
|
]
|
|
return chat_response_mock
|
|
|
|
mock_mistral_instance.chat.complete_async = mock_complete_async
|
|
|
|
messages: List[LLMMessage] = [
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
{"role": "user", "content": "What is async programming?"},
|
|
]
|
|
|
|
llm = MistralAILLM(model_name="mistral-model")
|
|
response = await llm.ainvoke(messages)
|
|
|
|
assert isinstance(response, LLMResponse)
|
|
assert response.content == "async mistral v2 response"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.SDKError", MockSDKError)
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
async def test_mistralai_llm_ainvoke_v2_error_handling(mock_mistral: Mock) -> None:
|
|
"""Test V2 interface async invoke method error handling."""
|
|
mock_mistral_instance = mock_mistral.return_value
|
|
|
|
async def mock_complete_async(*args: Any, **kwargs: Any) -> None:
|
|
raise MockSDKError("V2 async error")
|
|
|
|
mock_mistral_instance.chat.complete_async = mock_complete_async
|
|
|
|
messages: List[LLMMessage] = [
|
|
{"role": "user", "content": "This should fail"},
|
|
]
|
|
|
|
llm = MistralAILLM(model_name="mistral-model")
|
|
|
|
with pytest.raises(LLMGenerationError):
|
|
await llm.ainvoke(messages)
|
|
|
|
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
def test_mistralai_llm_invoke_v2_validation_error(mock_mistral: Mock) -> None:
|
|
"""Test V2 interface invoke with invalid message role raises error."""
|
|
mock_mistral_instance = mock_mistral.return_value
|
|
chat_response_mock = MagicMock()
|
|
chat_response_mock.choices = [
|
|
MagicMock(message=MagicMock(content="should not reach here"))
|
|
]
|
|
mock_mistral_instance.chat.complete.return_value = chat_response_mock
|
|
|
|
messages: List[LLMMessage] = [
|
|
{"role": "invalid_role", "content": "This should fail."}, # type: ignore[typeddict-item]
|
|
]
|
|
|
|
llm = MistralAILLM(model_name="mistral-model")
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
llm.invoke(messages)
|
|
assert "Unknown role: invalid_role" in str(exc_info.value)
|
|
|
|
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
def test_mistralai_llm_invoke_invalid_input_type(_mock_mistral: Mock) -> None:
|
|
"""Test that invalid input type raises appropriate error."""
|
|
llm = MistralAILLM(model_name="mistral-model")
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
llm.invoke(123) # type: ignore
|
|
assert "Invalid input type for invoke method" in str(exc_info.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
async def test_mistralai_llm_ainvoke_invalid_input_type(_mock_mistral: Mock) -> None:
|
|
"""Test that invalid input type raises appropriate error for async invoke."""
|
|
llm = MistralAILLM(model_name="mistral-model")
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
await llm.ainvoke(123) # type: ignore
|
|
assert "Invalid input type for ainvoke method" in str(exc_info.value)
|
|
|
|
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
def test_mistralai_llm_get_messages_v2_all_roles(_mock_mistral: Mock) -> None:
|
|
"""Test get_messages_v2 method handles all message roles correctly."""
|
|
messages: List[LLMMessage] = [
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi there!"},
|
|
{"role": "user", "content": "How are you?"},
|
|
]
|
|
|
|
llm = MistralAILLM(model_name="mistral-model")
|
|
result_messages = llm.get_messages_v2(messages)
|
|
|
|
# Verify the correct number of messages are returned
|
|
assert len(result_messages) == 4
|
|
|
|
# Verify each message type is correctly converted
|
|
assert result_messages[0].content == "You are a helpful assistant."
|
|
assert result_messages[1].content == "Hello"
|
|
assert result_messages[2].content == "Hi there!"
|
|
assert result_messages[3].content == "How are you?"
|
|
|
|
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
def test_mistralai_llm_get_messages_v2_unknown_role(_mock_mistral: Mock) -> None:
|
|
"""Test get_messages_v2 method raises error for unknown role."""
|
|
messages: List[LLMMessage] = [
|
|
{"role": "unknown_role", "content": "This should fail."}, # type: ignore[typeddict-item]
|
|
]
|
|
|
|
llm = MistralAILLM(model_name="mistral-model")
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
llm.get_messages_v2(messages)
|
|
assert "Unknown role: unknown_role" in str(exc_info.value)
|
|
|
|
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
def test_mistralai_invoke_v2_with_response_format_raises_error(
|
|
mock_mistral: Mock,
|
|
) -> None:
|
|
"""Test V2 interface raises NotImplementedError when response_format is used."""
|
|
|
|
class TestModel(BaseModel):
|
|
model_config = ConfigDict(extra="forbid")
|
|
value: str
|
|
|
|
messages: List[LLMMessage] = [{"role": "user", "content": "Test"}]
|
|
llm = MistralAILLM(api_key="test", model_name="mistral-model")
|
|
|
|
with pytest.raises(NotImplementedError) as exc_info:
|
|
llm.invoke(messages, response_format=TestModel)
|
|
|
|
assert "MistralAILLM does not currently support structured output" in str(
|
|
exc_info.value
|
|
)
|
|
|
|
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.SDKError", MockSDKError)
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
def test_mistralai_invoke_v2_rate_limit_handler_called(
|
|
mock_mistral: Mock,
|
|
) -> None:
|
|
"""Test that the rate limit handler is invoked on the V2 (List[LLMMessage]) path."""
|
|
messages: List[LLMMessage] = [{"role": "user", "content": "Hello"}]
|
|
mock_mistral_instance = mock_mistral.return_value
|
|
chat_response_mock = MagicMock()
|
|
chat_response_mock.choices = [MagicMock(message=MagicMock(content="Hi there!"))]
|
|
mock_mistral_instance.chat.complete.return_value = chat_response_mock
|
|
|
|
spy_handler = MagicMock(wraps=NoOpRateLimitHandler())
|
|
llm = MistralAILLM(model_name="mistral-model", rate_limit_handler=spy_handler)
|
|
response = llm.invoke(messages)
|
|
|
|
assert response.content == "Hi there!"
|
|
spy_handler.handle_sync.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.SDKError", MockSDKError)
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
async def test_mistralai_ainvoke_v2_rate_limit_handler_called(
|
|
mock_mistral: Mock,
|
|
) -> None:
|
|
"""Test that the rate limit handler is invoked on the async V2 (List[LLMMessage]) path."""
|
|
messages: List[LLMMessage] = [{"role": "user", "content": "Hello"}]
|
|
mock_mistral_instance = mock_mistral.return_value
|
|
chat_response_mock = MagicMock()
|
|
chat_response_mock.choices = [MagicMock(message=MagicMock(content="Hi there!"))]
|
|
mock_mistral_instance.chat.complete_async = AsyncMock(
|
|
return_value=chat_response_mock
|
|
)
|
|
|
|
spy_handler = MagicMock(wraps=NoOpRateLimitHandler())
|
|
llm = MistralAILLM(model_name="mistral-model", rate_limit_handler=spy_handler)
|
|
response = await llm.ainvoke(messages)
|
|
|
|
assert response.content == "Hi there!"
|
|
spy_handler.handle_async.assert_called_once()
|
|
|
|
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
def test_mistralai_llm_close(mock_mistral: Mock) -> None:
|
|
mock_mistral.return_value.aclose = AsyncMock()
|
|
|
|
llm = MistralAILLM(model_name="mistral-model")
|
|
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("error")
|
|
llm.close()
|
|
|
|
mock_mistral.return_value.close.assert_called_once()
|
|
mock_mistral.return_value.aclose.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("neo4j_graphrag.llm.mistralai_llm.Mistral")
|
|
async def test_mistralai_llm_aclose(mock_mistral: Mock) -> None:
|
|
mock_mistral.return_value.aclose = AsyncMock()
|
|
|
|
llm = MistralAILLM(model_name="mistral-model")
|
|
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("error")
|
|
await llm.aclose()
|
|
|
|
mock_mistral.return_value.close.assert_called_once()
|
|
mock_mistral.return_value.aclose.assert_called_once()
|