참고소스 수정본

참고/neo4j-graphrag-python-main/tests/unit/llm/conftest.py

@@ -0,0 +1,27 @@
+import pytest
+
+from neo4j_graphrag.tool import Tool, ObjectParameter, StringParameter
+
+
+class TestTool(Tool):
+    """Test tool for unit tests."""
+
+    def __init__(self, name: str = "test_tool", description: str = "A test tool"):
+        parameters = ObjectParameter(
+            description="Test parameters",
+            properties={"param1": StringParameter(description="Test parameter")},
+            required_properties=["param1"],
+            additional_properties=False,
+        )
+
+        super().__init__(
+            name=name,
+            description=description,
+            parameters=parameters,
+            execute_func=lambda **kwargs: kwargs,
+        )
+
+
+@pytest.fixture
+def test_tool() -> Tool:
+    return TestTool()

참고/neo4j-graphrag-python-main/tests/unit/llm/test_anthropic_llm.py

@@ -0,0 +1,460 @@
+#  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.
+from __future__ import annotations
+import warnings
+import sys
+from typing import Any, Generator, List, cast
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
+
+import anthropic
+import pytest
+from neo4j_graphrag.exceptions import LLMGenerationError
+from neo4j_graphrag.llm.anthropic_llm import AnthropicLLM
+from neo4j_graphrag.llm.types import LLMResponse
+from neo4j_graphrag.types import LLMMessage
+from pydantic import BaseModel, ConfigDict
+
+
+@pytest.fixture
+def mock_anthropic() -> Generator[MagicMock, None, None]:
+    mock = MagicMock()
+    mock.APIError = anthropic.APIError
+    mock.NOT_GIVEN = anthropic.NOT_GIVEN
+
+    with patch.dict(sys.modules, {"anthropic": mock}):
+        yield mock
+
+
+def _as_mock(value: Any) -> MagicMock:
+    return cast(MagicMock, value)
+
+
+def _as_async_mock(value: Any) -> AsyncMock:
+    return cast(AsyncMock, value)
+
+
+@patch("builtins.__import__", side_effect=ImportError)
+def test_anthropic_llm_missing_dependency(mock_import: Mock) -> None:
+    with pytest.raises(ImportError):
+        AnthropicLLM(model_name="claude-3-opus-20240229")
+
+
+def test_anthropic_invoke_happy_path(mock_anthropic: Mock) -> None:
+    mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
+        content=[MagicMock(text="generated text")]
+    )
+    model_params = {"temperature": 0.3}
+    llm = AnthropicLLM("claude-3-opus-20240229", model_params=model_params)
+    input_text = "may thy knife chip and shatter"
+    response = llm.invoke(input_text)
+    assert response.content == "generated text"
+    _as_mock(llm.client.messages.create).assert_called_once_with(
+        messages=[{"role": "user", "content": input_text}],
+        model="claude-3-opus-20240229",
+        system=anthropic.NOT_GIVEN,
+        **model_params,
+    )
+
+
+def test_anthropic_invoke_with_message_history_happy_path(mock_anthropic: Mock) -> None:
+    mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
+        content=[MagicMock(text="generated text")]
+    )
+    model_params = {"temperature": 0.3}
+    llm = AnthropicLLM(
+        "claude-3-opus-20240229",
+        model_params=model_params,
+    )
+    message_history = [
+        {"role": "user", "content": "When does the sun come up in the summer?"},
+        {"role": "assistant", "content": "Usually around 6am."},
+    ]
+    question = "What about next season?"
+
+    response = llm.invoke(question, message_history)  # type: ignore
+    assert response.content == "generated text"
+    message_history.append({"role": "user", "content": question})
+    _as_mock(llm.client.messages.create).assert_called_once_with(
+        messages=message_history,
+        model="claude-3-opus-20240229",
+        system=anthropic.NOT_GIVEN,
+        **model_params,
+    )
+
+
+def test_anthropic_invoke_with_system_instruction(
+    mock_anthropic: Mock,
+) -> None:
+    mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
+        content=[MagicMock(text="generated text")]
+    )
+    model_params = {"temperature": 0.3}
+    system_instruction = "You are a helpful assistant."
+    llm = AnthropicLLM(
+        "claude-3-opus-20240229",
+        model_params=model_params,
+    )
+
+    question = "When does it come up in the winter?"
+    response = llm.invoke(question, system_instruction=system_instruction)
+    assert isinstance(response, LLMResponse)
+    assert response.content == "generated text"
+    messages: List[LLMMessage] = [{"role": "user", "content": question}]
+    _as_mock(llm.client.messages.create).assert_called_with(
+        model="claude-3-opus-20240229",
+        system=system_instruction,
+        messages=messages,
+        **model_params,
+    )
+
+    assert _as_mock(llm.client.messages.create).call_count == 1
+
+
+def test_anthropic_invoke_with_message_history_and_system_instruction(
+    mock_anthropic: Mock,
+) -> None:
+    mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
+        content=[MagicMock(text="generated text")]
+    )
+    model_params = {"temperature": 0.3}
+    system_instruction = "You are a helpful assistant."
+    llm = AnthropicLLM(
+        "claude-3-opus-20240229",
+        model_params=model_params,
+    )
+    message_history = [
+        {"role": "user", "content": "When does the sun come up in the summer?"},
+        {"role": "assistant", "content": "Usually around 6am."},
+    ]
+
+    question = "When does it come up in the winter?"
+    response = llm.invoke(question, message_history, system_instruction)  # type: ignore
+    assert isinstance(response, LLMResponse)
+    assert response.content == "generated text"
+    message_history.append({"role": "user", "content": question})
+    _as_mock(llm.client.messages.create).assert_called_with(
+        model="claude-3-opus-20240229",
+        system=system_instruction,
+        messages=message_history,
+        **model_params,
+    )
+
+    assert _as_mock(llm.client.messages.create).call_count == 1
+
+
+def test_anthropic_invoke_with_message_history_validation_error(
+    mock_anthropic: Mock,
+) -> None:
+    mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
+        content=[MagicMock(text="generated text")]
+    )
+    model_params = {"temperature": 0.3}
+    system_instruction = "You are a helpful assistant."
+    llm = AnthropicLLM(
+        "claude-3-opus-20240229",
+        model_params=model_params,
+        system_instruction=system_instruction,
+    )
+    message_history = [
+        {"role": "human", "content": "When does the sun come up in the summer?"},
+        {"role": "assistant", "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
+async def test_anthropic_ainvoke_happy_path(mock_anthropic: Mock) -> None:
+    mock_response = AsyncMock()
+    mock_response.content = [MagicMock(text="Return text")]
+    mock_model = mock_anthropic.AsyncAnthropic.return_value
+    mock_model.messages.create = AsyncMock(return_value=mock_response)
+    model_params = {"temperature": 0.3}
+    llm = AnthropicLLM("claude-3-opus-20240229", model_params)
+    input_text = "may thy knife chip and shatter"
+    response = await llm.ainvoke(input_text)
+    assert response.content == "Return text"
+    _as_async_mock(llm.async_client.messages.create).assert_awaited_once_with(
+        model="claude-3-opus-20240229",
+        system=anthropic.NOT_GIVEN,
+        messages=[{"role": "user", "content": input_text}],
+        **model_params,
+    )
+
+
+# V2 Interface Tests
+
+
+def test_anthropic_llm_invoke_v2_happy_path(mock_anthropic: Mock) -> None:
+    """Test V2 interface invoke method with List[LLMMessage] input."""
+    mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
+        content=[MagicMock(text="anthropic v2 response")]
+    )
+    mock_anthropic.types.MessageParam = MagicMock(side_effect=lambda **kwargs: kwargs)
+
+    messages: List[LLMMessage] = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "What is machine learning?"},
+    ]
+
+    model_params = {"temperature": 0.7}
+    llm = AnthropicLLM(model_name="claude-3-opus-20240229", model_params=model_params)
+    response = llm.invoke(messages)
+
+    assert isinstance(response, LLMResponse)
+    assert response.content == "anthropic v2 response"
+
+    # Verify the correct method was called with system instruction and messages
+    _as_mock(llm.client.messages.create).assert_called_once_with(
+        model="claude-3-opus-20240229",
+        system="You are a helpful assistant.",
+        messages=[{"role": "user", "content": "What is machine learning?"}],
+        **model_params,
+    )
+
+
+def test_anthropic_llm_invoke_v2_with_conversation_history(
+    mock_anthropic: Mock,
+) -> None:
+    """Test V2 interface invoke method with complex conversation history."""
+    mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
+        content=[MagicMock(text="anthropic conversation response")]
+    )
+    mock_anthropic.types.MessageParam = MagicMock(side_effect=lambda **kwargs: kwargs)
+
+    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 = AnthropicLLM(model_name="claude-3-opus-20240229")
+    response = llm.invoke(messages)
+
+    assert isinstance(response, LLMResponse)
+    assert response.content == "anthropic conversation response"
+
+    # Verify the correct number of messages were passed (excluding system)
+    _as_mock(llm.client.messages.create).assert_called_once()
+    call_args = _as_mock(llm.client.messages.create).call_args[1]
+    assert call_args["system"] == "You are a helpful assistant."
+    assert len(call_args["messages"]) == 3
+
+
+def test_anthropic_llm_invoke_v2_no_system_message(mock_anthropic: Mock) -> None:
+    """Test V2 interface invoke method without system message."""
+    mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
+        content=[MagicMock(text="anthropic no system response")]
+    )
+    mock_anthropic.types.MessageParam = MagicMock(side_effect=lambda **kwargs: kwargs)
+
+    messages: List[LLMMessage] = [
+        {"role": "user", "content": "What is the capital of France?"},
+    ]
+
+    llm = AnthropicLLM(model_name="claude-3-opus-20240229")
+    response = llm.invoke(messages)
+
+    assert isinstance(response, LLMResponse)
+    assert response.content == "anthropic no system response"
+
+    # Verify only user message was passed and no system instruction
+    _as_mock(llm.client.messages.create).assert_called_once()
+    call_args = _as_mock(llm.client.messages.create).call_args[1]
+    assert call_args["system"] == anthropic.NOT_GIVEN
+    assert len(call_args["messages"]) == 1
+
+
+@pytest.mark.asyncio
+async def test_anthropic_llm_ainvoke_v2_happy_path(mock_anthropic: Mock) -> None:
+    """Test V2 interface async invoke method with List[LLMMessage] input."""
+    mock_response = AsyncMock()
+    mock_response.content = [MagicMock(text="async anthropic v2 response")]
+    mock_model = mock_anthropic.AsyncAnthropic.return_value
+    mock_model.messages.create = AsyncMock(return_value=mock_response)
+    mock_anthropic.types.MessageParam = MagicMock(side_effect=lambda **kwargs: kwargs)
+
+    messages: List[LLMMessage] = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "What is async programming?"},
+    ]
+
+    model_params = {"max_tokens": 100}
+    llm = AnthropicLLM(model_name="claude-3-opus-20240229", model_params=model_params)
+    response = await llm.ainvoke(messages)
+
+    assert isinstance(response, LLMResponse)
+    assert response.content == "async anthropic v2 response"
+
+    # Verify the async client was called correctly
+    _as_async_mock(llm.async_client.messages.create).assert_awaited_once_with(
+        model="claude-3-opus-20240229",
+        system="You are a helpful assistant.",
+        messages=[{"role": "user", "content": "What is async programming?"}],
+        **model_params,
+    )
+
+
+def test_anthropic_llm_invoke_v2_validation_error(mock_anthropic: Mock) -> None:
+    """Test V2 interface invoke method with invalid role."""
+    mock_anthropic.types.MessageParam = MagicMock(side_effect=lambda **kwargs: kwargs)
+
+    messages: List[LLMMessage] = [
+        {"role": "invalid_role", "content": "This should fail."},  # type: ignore[typeddict-item]
+    ]
+
+    llm = AnthropicLLM(model_name="claude-3-opus-20240229")
+
+    with pytest.raises(ValueError) as exc_info:
+        llm.invoke(messages)
+    assert "Unknown role: invalid_role" in str(exc_info.value)
+
+
+def test_anthropic_llm_invoke_invalid_input_type(
+    mock_anthropic: Mock,
+) -> None:  # noqa: ARG001
+    """Test that invalid input type raises appropriate error."""
+    llm = AnthropicLLM(model_name="claude-3-opus-20240229")
+
+    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
+async def test_anthropic_llm_ainvoke_invalid_input_type(
+    mock_anthropic: Mock,
+) -> None:  # noqa: ARG001
+    """Test that invalid input type raises appropriate error for async invoke."""
+    llm = AnthropicLLM(model_name="claude-3-opus-20240229")
+
+    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)
+
+
+def test_anthropic_llm_get_brand_new_messages_all_roles(mock_anthropic: Mock) -> None:
+    """Test get_brand_new_messages method handles all message roles correctly."""
+
+    def create_message_param(**kwargs: str) -> MagicMock:
+        return MagicMock(**kwargs)
+
+    mock_anthropic.types.MessageParam = MagicMock(side_effect=create_message_param)
+
+    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 = AnthropicLLM(model_name="claude-3-opus-20240229")
+    system_instruction, result_messages = llm.get_messages_v2(messages)
+
+    # Verify system instruction is extracted
+    assert system_instruction == "You are a helpful assistant."
+
+    result_messages = cast(list[MagicMock], list(result_messages))
+
+    # Verify the correct number of non-system messages are returned
+    assert len(result_messages) == 3
+
+    # Verify message content is preserved
+    assert result_messages[0].content == "Hello"
+    assert result_messages[1].content == "Hi there!"
+    assert result_messages[2].content == "How are you?"
+
+
+def test_anthropic_llm_get_brand_new_messages_unknown_role(
+    mock_anthropic: Mock,
+) -> None:  # noqa: ARG001
+    """Test get_brand_new_messages method raises error for unknown role."""
+    messages: List[LLMMessage] = [
+        {"role": "unknown_role", "content": "This should fail."},  # type: ignore[typeddict-item]
+    ]
+
+    llm = AnthropicLLM(model_name="claude-3-opus-20240229")
+
+    with pytest.raises(ValueError) as exc_info:
+        llm.get_messages_v2(messages)
+    assert "Unknown role: unknown_role" in str(exc_info.value)
+
+
+def test_anthropic_llm_invoke_v2_empty_response_error(mock_anthropic: Mock) -> None:
+    """Test V2 interface invoke method handles empty response."""
+    mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
+        content=[]  # Empty content should trigger error
+    )
+    mock_anthropic.types.MessageParam = MagicMock(side_effect=lambda **kwargs: kwargs)
+
+    messages: List[LLMMessage] = [
+        {"role": "user", "content": "This should return empty response."},
+    ]
+
+    llm = AnthropicLLM(model_name="claude-3-opus-20240229")
+
+    with pytest.raises(LLMGenerationError) as exc_info:
+        llm.invoke(messages)
+    assert "LLM returned empty response" in str(exc_info.value)
+
+
+def test_anthropic_invoke_v2_with_response_format_raises_error(
+    mock_anthropic: 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 = AnthropicLLM(api_key="test", model_name="claude-3-opus")
+
+    with pytest.raises(NotImplementedError) as exc_info:
+        llm.invoke(messages, response_format=TestModel)
+
+    assert "AnthropicLLM does not currently support structured output" in str(
+        exc_info.value
+    )
+
+
+def test_anthropic_llm_close(mock_anthropic: Mock) -> None:
+    mock_anthropic.AsyncAnthropic.return_value.close = AsyncMock()
+
+    llm = AnthropicLLM("claude-3-opus-20240229")
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        llm.close()
+
+    mock_anthropic.Anthropic.return_value.close.assert_called_once()
+    mock_anthropic.AsyncAnthropic.return_value.close.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_anthropic_llm_aclose(mock_anthropic: Mock) -> None:
+    mock_anthropic.AsyncAnthropic.return_value.close = AsyncMock()
+
+    llm = AnthropicLLM("claude-3-opus-20240229")
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        await llm.aclose()
+
+    mock_anthropic.Anthropic.return_value.close.assert_called_once()
+    mock_anthropic.AsyncAnthropic.return_value.close.assert_called_once()

참고/neo4j-graphrag-python-main/tests/unit/llm/test_base_llm.py

@@ -0,0 +1,211 @@
+#  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, Type, Union
+
+import pytest
+from pydantic import BaseModel, ValidationError
+
+from neo4j_graphrag.llm.base import LLMBase
+from neo4j_graphrag.llm.types import LLMResponse, LLMUsage
+from neo4j_graphrag.message_history import MessageHistory
+from neo4j_graphrag.types import LLMMessage
+from neo4j_graphrag.utils.rate_limit import NoOpRateLimitHandler
+
+# ---------------------------------------------------------------------------
+# LLMUsage
+# ---------------------------------------------------------------------------
+
+
+def test_llm_usage_defaults_to_none() -> None:
+    usage = LLMUsage()
+    assert usage.request_tokens is None
+    assert usage.response_tokens is None
+    assert usage.total_tokens is None
+
+
+def test_llm_usage_accepts_explicit_values() -> None:
+    usage = LLMUsage(request_tokens=10, response_tokens=20, total_tokens=30)
+    assert usage.request_tokens == 10
+    assert usage.response_tokens == 20
+    assert usage.total_tokens == 30
+
+
+def test_llm_usage_partial_values_keep_other_defaults() -> None:
+    usage = LLMUsage(request_tokens=5)
+    assert usage.request_tokens == 5
+    assert usage.response_tokens is None
+    assert usage.total_tokens is None
+
+
+def test_llm_usage_rejects_non_integer_tokens() -> None:
+    with pytest.raises(ValidationError):
+        LLMUsage(request_tokens="bad")  # type: ignore[arg-type]
+
+
+def test_llm_response_usage_is_none_by_default() -> None:
+    response = LLMResponse(content="hello")
+    assert response.usage is None
+
+
+def test_llm_response_carries_usage() -> None:
+    usage = LLMUsage(request_tokens=3, response_tokens=7, total_tokens=10)
+    response = LLMResponse(content="hi", usage=usage)
+    assert response.usage is not None
+    assert response.usage.request_tokens == 3
+    assert response.usage.response_tokens == 7
+    assert response.usage.total_tokens == 10
+
+
+# ---------------------------------------------------------------------------
+# Minimal concrete subclass used across tests
+# ---------------------------------------------------------------------------
+
+
+class _ConcreteLLM(LLMBase):
+    """Minimal LLMBase subclass for unit testing."""
+
+    def invoke(
+        self,
+        input: Union[str, List[LLMMessage]],
+        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
+        system_instruction: Optional[str] = None,
+        response_format: Optional[Union[Type[BaseModel], dict[str, Any]]] = None,
+        **kwargs: Any,
+    ) -> LLMResponse:
+        if isinstance(input, str):
+            return LLMResponse(content=f"v1:{input}")
+        return LLMResponse(content="v2:list")
+
+    async def ainvoke(
+        self,
+        input: Union[str, List[LLMMessage]],
+        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
+        system_instruction: Optional[str] = None,
+        response_format: Optional[Union[Type[BaseModel], dict[str, Any]]] = None,
+        **kwargs: Any,
+    ) -> LLMResponse:
+        if isinstance(input, str):
+            return LLMResponse(content=f"async_v1:{input}")
+        return LLMResponse(content="async_v2:list")
+
+
+# ---------------------------------------------------------------------------
+# Instantiation
+# ---------------------------------------------------------------------------
+
+
+def test_llmbase_cannot_be_instantiated_directly() -> None:
+    with pytest.raises(TypeError):
+        LLMBase(model_name="m")
+
+
+def test_llmbase_sets_model_name() -> None:
+    llm = _ConcreteLLM(model_name="my-model")
+    assert llm.model_name == "my-model"
+
+
+def test_llmbase_default_model_params_is_empty_dict() -> None:
+    llm = _ConcreteLLM(model_name="m")
+    assert llm.model_params == {}
+
+
+def test_llmbase_accepts_model_params() -> None:
+    llm = _ConcreteLLM(model_name="m", model_params={"temperature": 0.5})
+    assert llm.model_params == {"temperature": 0.5}
+
+
+def test_llmbase_accepts_custom_rate_limit_handler() -> None:
+    handler = NoOpRateLimitHandler()
+    llm = _ConcreteLLM(model_name="m", rate_limit_handler=handler)
+    assert llm._rate_limit_handler is handler
+
+
+def test_llmbase_init_does_not_emit_deprecation_warning() -> None:
+    """LLMBase.__init__ delegates to LLMInterfaceV2, which has no deprecation warning."""
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        _ConcreteLLM(model_name="m")
+    deprecation_warnings = [
+        w for w in caught if issubclass(w.category, DeprecationWarning)
+    ]
+    assert deprecation_warnings == []
+
+
+# ---------------------------------------------------------------------------
+# invoke routing
+# ---------------------------------------------------------------------------
+
+
+def test_invoke_with_str_routes_to_v1() -> None:
+    llm = _ConcreteLLM(model_name="m")
+    result = llm.invoke("hello")
+    assert result.content == "v1:hello"
+
+
+def test_invoke_with_list_routes_to_v2() -> None:
+    llm = _ConcreteLLM(model_name="m")
+    messages: List[LLMMessage] = [{"role": "user", "content": "hi"}]
+    result = llm.invoke(messages)
+    assert result.content == "v2:list"
+
+
+def test_invoke_v2_accepts_response_format_kwarg() -> None:
+    class MyModel(BaseModel):
+        answer: str
+
+    llm = _ConcreteLLM(model_name="m")
+    messages: List[LLMMessage] = [{"role": "user", "content": "hi"}]
+    # response_format must be keyword-only; this should not raise
+    result = llm.invoke(messages, response_format=MyModel)
+    assert result.content == "v2:list"
+
+
+# ---------------------------------------------------------------------------
+# ainvoke routing
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_ainvoke_with_str_routes_to_v1() -> None:
+    llm = _ConcreteLLM(model_name="m")
+    result = await llm.ainvoke("hello")
+    assert result.content == "async_v1:hello"
+
+
+@pytest.mark.asyncio
+async def test_ainvoke_with_list_routes_to_v2() -> None:
+    llm = _ConcreteLLM(model_name="m")
+    messages: List[LLMMessage] = [{"role": "user", "content": "hi"}]
+    result = await llm.ainvoke(messages)
+    assert result.content == "async_v2:list"
+
+
+# ---------------------------------------------------------------------------
+# Tool calling defaults (inherited from LLMInterface)
+# ---------------------------------------------------------------------------
+
+
+def test_invoke_with_tools_raises_not_implemented() -> None:
+    llm = _ConcreteLLM(model_name="m")
+    with pytest.raises(NotImplementedError):
+        llm.invoke_with_tools("hello", tools=[])
+
+
+@pytest.mark.asyncio
+async def test_ainvoke_with_tools_raises_not_implemented() -> None:
+    llm = _ConcreteLLM(model_name="m")
+    with pytest.raises(NotImplementedError):
+        await llm.ainvoke_with_tools("hello", tools=[])

참고/neo4j-graphrag-python-main/tests/unit/llm/test_bedrock_llm.py

@@ -0,0 +1,198 @@
+#  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.
+from __future__ import annotations
+
+from typing import Any, Generator
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from neo4j_graphrag.exceptions import LLMGenerationError
+from neo4j_graphrag.llm import BedrockLLM
+from neo4j_graphrag.types import LLMMessage
+
+
+@pytest.fixture
+def mock_boto3() -> Generator[MagicMock, None, None]:
+    with patch("neo4j_graphrag.llm.bedrock_llm.boto3") as mock_boto:
+        mock_client = MagicMock()
+        mock_boto.client.return_value = mock_client
+        yield mock_boto
+
+
+def _make_converse_response(text: str = "generated text") -> dict[str, Any]:
+    return {
+        "output": {
+            "message": {
+                "role": "assistant",
+                "content": [{"text": text}],
+            }
+        }
+    }
+
+
+def test_bedrock_llm_missing_dependency() -> None:
+    with patch("neo4j_graphrag.llm.bedrock_llm.boto3", None):
+        with pytest.raises(ImportError) as exc:
+            BedrockLLM(model_name="us.anthropic.claude-sonnet-4-20250514-v1:0")
+        assert "Could not import boto3 python client" in str(exc.value)
+
+
+def test_bedrock_llm_default_model_from_env(mock_boto3: MagicMock) -> None:
+    with patch.dict("os.environ", {"BEDROCK_LLM_MODEL": "custom-llm-model"}):
+        import importlib
+        import sys
+
+        original_boto3 = sys.modules.get("boto3")
+        sys.modules["boto3"] = mock_boto3
+
+        try:
+            import neo4j_graphrag.llm.bedrock_llm as bedrock_llm_mod
+
+            importlib.reload(bedrock_llm_mod)
+
+            assert bedrock_llm_mod.DEFAULT_BEDROCK_LLM_MODEL == "custom-llm-model"
+
+            llm = bedrock_llm_mod.BedrockLLM()
+            assert llm.model_name == "custom-llm-model"
+        finally:
+            if original_boto3 is not None:
+                sys.modules["boto3"] = original_boto3
+            importlib.reload(bedrock_llm_mod)
+
+
+def test_bedrock_invoke_happy_path(mock_boto3: MagicMock) -> None:
+    mock_client = mock_boto3.client.return_value
+    mock_client.converse.return_value = _make_converse_response("hello world")
+
+    llm = BedrockLLM("us.anthropic.claude-sonnet-4-20250514-v1:0")
+    response = llm.invoke("hello")
+
+    assert response.content == "hello world"
+    mock_client.converse.assert_called_once()
+
+
+def test_bedrock_invoke_with_message_history(mock_boto3: MagicMock) -> None:
+    mock_client = mock_boto3.client.return_value
+    mock_client.converse.return_value = _make_converse_response("response")
+
+    llm = BedrockLLM("us.anthropic.claude-sonnet-4-20250514-v1:0")
+    history: list[LLMMessage] = [
+        {"role": "user", "content": "previous question"},
+        {"role": "assistant", "content": "previous answer"},
+    ]
+    response = llm.invoke("follow up", message_history=history)
+
+    assert response.content == "response"
+    call_kwargs = mock_client.converse.call_args[1]
+    # 2 history messages + 1 new user message
+    assert len(call_kwargs["messages"]) == 3
+
+
+def test_bedrock_invoke_with_system_instruction(mock_boto3: MagicMock) -> None:
+    mock_client = mock_boto3.client.return_value
+    mock_client.converse.return_value = _make_converse_response("response")
+
+    llm = BedrockLLM("us.anthropic.claude-sonnet-4-20250514-v1:0")
+    response = llm.invoke("hello", system_instruction="You are a bot")
+
+    assert response.content == "response"
+    call_kwargs = mock_client.converse.call_args[1]
+    assert call_kwargs["system"] == [{"text": "You are a bot"}]
+
+
+@pytest.mark.asyncio
+async def test_bedrock_ainvoke_happy_path(mock_boto3: MagicMock) -> None:
+    mock_client = mock_boto3.client.return_value
+    mock_client.converse.return_value = _make_converse_response("async response")
+
+    llm = BedrockLLM("us.anthropic.claude-sonnet-4-20250514-v1:0")
+    response = await llm.ainvoke("hello")
+
+    assert response.content == "async response"
+    mock_client.converse.assert_called_once()
+
+
+def test_bedrock_invoke_v2_happy_path(mock_boto3: MagicMock) -> None:
+    mock_client = mock_boto3.client.return_value
+    mock_client.converse.return_value = _make_converse_response("v2 response")
+
+    messages: list[LLMMessage] = [
+        {"role": "system", "content": "You are a bot"},
+        {"role": "user", "content": "hello"},
+    ]
+
+    llm = BedrockLLM("us.anthropic.claude-sonnet-4-20250514-v1:0")
+    response = llm.invoke(messages)
+
+    assert response.content == "v2 response"
+    call_kwargs = mock_client.converse.call_args[1]
+    assert call_kwargs["system"] == [{"text": "You are a bot"}]
+    # only user message, system is extracted
+    assert len(call_kwargs["messages"]) == 1
+
+
+@pytest.mark.asyncio
+async def test_bedrock_ainvoke_v2_happy_path(mock_boto3: MagicMock) -> None:
+    mock_client = mock_boto3.client.return_value
+    mock_client.converse.return_value = _make_converse_response("async v2")
+
+    messages: list[LLMMessage] = [{"role": "user", "content": "hello"}]
+
+    llm = BedrockLLM("us.anthropic.claude-sonnet-4-20250514-v1:0")
+    response = await llm.ainvoke(messages)
+
+    assert response.content == "async v2"
+
+
+def test_bedrock_invoke_error(mock_boto3: MagicMock) -> None:
+    mock_client = mock_boto3.client.return_value
+    mock_client.converse.side_effect = Exception("API error")
+
+    llm = BedrockLLM("us.anthropic.claude-sonnet-4-20250514-v1:0")
+    with pytest.raises(LLMGenerationError):
+        llm.invoke("hello")
+
+
+def test_bedrock_invoke_empty_response(mock_boto3: MagicMock) -> None:
+    mock_client = mock_boto3.client.return_value
+    mock_client.converse.return_value = {"output": {"message": {"content": []}}}
+
+    llm = BedrockLLM("us.anthropic.claude-sonnet-4-20250514-v1:0")
+    with pytest.raises(LLMGenerationError, match="LLM returned empty response"):
+        llm.invoke("hello")
+
+
+def test_bedrock_invoke_v2_with_response_format_raises_error(
+    mock_boto3: MagicMock,
+) -> None:
+    messages: list[LLMMessage] = [{"role": "user", "content": "hello"}]
+    llm = BedrockLLM("us.anthropic.claude-sonnet-4-20250514-v1:0")
+    with pytest.raises(NotImplementedError):
+        llm.invoke(messages, response_format={"type": "json_object"})
+
+
+def test_bedrock_invoke_with_model_params(mock_boto3: MagicMock) -> None:
+    mock_client = mock_boto3.client.return_value
+    mock_client.converse.return_value = _make_converse_response("response")
+
+    llm = BedrockLLM(
+        "us.anthropic.claude-sonnet-4-20250514-v1:0",
+        model_params={"temperature": 0.5, "maxTokens": 512},
+    )
+    llm.invoke("hello")
+
+    call_kwargs = mock_client.converse.call_args[1]
+    assert call_kwargs["inferenceConfig"] == {"temperature": 0.5, "maxTokens": 512}

참고/neo4j-graphrag-python-main/tests/unit/llm/test_cohere_llm.py

@@ -0,0 +1,303 @@
+#  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 sys
+import warnings
+from typing import Generator, List
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
+
+import cohere.core
+import pytest
+from neo4j_graphrag.exceptions import LLMGenerationError
+from neo4j_graphrag.llm import LLMResponse
+from neo4j_graphrag.llm.cohere_llm import CohereLLM
+from neo4j_graphrag.types import LLMMessage
+from pydantic import BaseModel, ConfigDict
+
+
+@pytest.fixture
+def mock_cohere() -> Generator[MagicMock, None, None]:
+    mock_cohere = MagicMock()
+    mock_cohere.core.api_error.ApiError = cohere.core.ApiError
+    with patch.dict(sys.modules, {"cohere": mock_cohere}):
+        yield mock_cohere
+
+
+@patch("builtins.__import__", side_effect=ImportError)
+def test_cohere_llm_missing_dependency(mock_import: Mock) -> None:
+    with pytest.raises(ImportError):
+        CohereLLM(model_name="something")
+
+
+def test_cohere_llm_happy_path(mock_cohere: Mock) -> None:
+    chat_response_mock = MagicMock()
+    chat_response_mock.message.content = [MagicMock(text="cohere response text")]
+    mock_cohere.ClientV2.return_value.chat.return_value = chat_response_mock
+    llm = CohereLLM(model_name="something")
+    res = llm.invoke("my text")
+    assert isinstance(res, LLMResponse)
+    assert res.content == "cohere response text"
+
+
+def test_cohere_llm_invoke_with_message_history_happy_path(mock_cohere: Mock) -> None:
+    chat_response_mock = MagicMock()
+    chat_response_mock.message.content = [MagicMock(text="cohere response text")]
+    mock_cohere_client_chat = mock_cohere.ClientV2.return_value.chat
+    mock_cohere_client_chat.return_value = chat_response_mock
+
+    system_instruction = "You are a helpful assistant."
+    llm = CohereLLM(model_name="something")
+    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 == "cohere response text"
+    messages: List[LLMMessage] = [{"role": "system", "content": system_instruction}]
+    messages.extend(message_history)
+    messages.append({"role": "user", "content": question})
+    mock_cohere_client_chat.assert_called_once_with(
+        messages=messages,
+        model="something",
+    )
+
+
+def test_cohere_llm_invoke_with_message_history_and_system_instruction(
+    mock_cohere: Mock,
+) -> None:
+    chat_response_mock = MagicMock()
+    chat_response_mock.message.content = [MagicMock(text="cohere response text")]
+    mock_cohere_client_chat = mock_cohere.ClientV2.return_value.chat
+    mock_cohere_client_chat.return_value = chat_response_mock
+
+    system_instruction = "You are a helpful assistant."
+    llm = CohereLLM(model_name="gpt")
+    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 == "cohere response text"
+    messages: List[LLMMessage] = [{"role": "system", "content": system_instruction}]
+    messages.extend(message_history)
+    messages.append({"role": "user", "content": question})
+    mock_cohere_client_chat.assert_called_once_with(
+        messages=messages,
+        model="gpt",
+    )
+
+
+def test_cohere_llm_invoke_with_message_history_validation_error(
+    mock_cohere: Mock,
+) -> None:
+    chat_response_mock = MagicMock()
+    chat_response_mock.message.content = [MagicMock(text="cohere response text")]
+    mock_cohere.ClientV2.return_value.chat.return_value = chat_response_mock
+
+    system_instruction = "You are a helpful assistant."
+    llm = CohereLLM(model_name="something", system_instruction=system_instruction)
+    message_history = [
+        {"role": "robot", "content": "When does the sun come up in the summer?"},
+        {"role": "assistant", "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
+async def test_cohere_llm_happy_path_async(mock_cohere: Mock) -> None:
+    chat_response_mock = MagicMock(
+        message=MagicMock(content=[MagicMock(text="cohere response text")])
+    )
+    mock_cohere.AsyncClientV2.return_value.chat = AsyncMock(
+        return_value=chat_response_mock
+    )
+
+    llm = CohereLLM(model_name="something")
+    res = await llm.ainvoke("my text")
+    assert isinstance(res, LLMResponse)
+    assert res.content == "cohere response text"
+
+
+def test_cohere_llm_failed(mock_cohere: Mock) -> None:
+    mock_cohere.ClientV2.return_value.chat.side_effect = cohere.core.ApiError
+    llm = CohereLLM(model_name="something")
+    with pytest.raises(LLMGenerationError) as excinfo:
+        llm.invoke("my text")
+    assert "ApiError" in str(excinfo)
+
+
+@pytest.mark.asyncio
+async def test_cohere_llm_failed_async(mock_cohere: Mock) -> None:
+    mock_cohere.AsyncClientV2.return_value.chat.side_effect = cohere.core.ApiError
+    llm = CohereLLM(model_name="something")
+
+    with pytest.raises(LLMGenerationError) as excinfo:
+        await llm.ainvoke("my text")
+    assert "ApiError" in str(excinfo)
+
+
+# V2 Interface Tests
+
+
+def test_cohere_llm_invoke_v2_happy_path(mock_cohere: Mock) -> None:
+    """Test V2 interface invoke method with List[LLMMessage] input."""
+    chat_response_mock = MagicMock()
+    chat_response_mock.message.content = [MagicMock(text="cohere v2 response text")]
+    mock_cohere.ClientV2.return_value.chat.return_value = chat_response_mock
+
+    # Mock Cohere message types
+    mock_cohere.SystemChatMessageV2 = MagicMock()
+    mock_cohere.UserChatMessageV2 = MagicMock()
+    mock_cohere.AssistantChatMessageV2 = MagicMock()
+
+    messages: List[LLMMessage] = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "What is the capital of France?"},
+    ]
+
+    llm = CohereLLM(model_name="something")
+    response = llm.invoke(messages)
+
+    assert isinstance(response, LLMResponse)
+    assert response.content == "cohere v2 response text"
+
+    # Verify the client was called correctly
+    mock_cohere.ClientV2.return_value.chat.assert_called_once()
+    call_args = mock_cohere.ClientV2.return_value.chat.call_args[1]
+    assert call_args["model"] == "something"
+
+
+@pytest.mark.asyncio
+async def test_cohere_llm_ainvoke_v2_happy_path(mock_cohere: Mock) -> None:
+    """Test V2 interface async invoke method with List[LLMMessage] input."""
+    chat_response_mock = MagicMock()
+    chat_response_mock.message.content = [
+        MagicMock(text="cohere v2 async response text")
+    ]
+    mock_cohere.AsyncClientV2.return_value.chat = AsyncMock(
+        return_value=chat_response_mock
+    )
+
+    # Mock Cohere message types
+    mock_cohere.SystemChatMessageV2 = MagicMock()
+    mock_cohere.UserChatMessageV2 = MagicMock()
+    mock_cohere.AssistantChatMessageV2 = MagicMock()
+
+    messages: List[LLMMessage] = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "What is the capital of France?"},
+    ]
+
+    llm = CohereLLM(model_name="something")
+    response = await llm.ainvoke(messages)
+
+    assert isinstance(response, LLMResponse)
+    assert response.content == "cohere v2 async response text"
+
+    # Verify the async client was called correctly
+    mock_cohere.AsyncClientV2.return_value.chat.assert_awaited_once()
+
+
+def test_cohere_llm_invoke_v2_validation_error(mock_cohere: Mock) -> None:
+    """Test V2 interface invoke with invalid message role raises error."""
+    chat_response_mock = MagicMock()
+    chat_response_mock.message.content = [MagicMock(text="should not get here")]
+    mock_cohere.ClientV2.return_value.chat.return_value = chat_response_mock
+
+    messages: List[LLMMessage] = [
+        {"role": "invalid_role", "content": "This should fail."},  # type: ignore[typeddict-item]
+    ]
+
+    llm = CohereLLM(model_name="something")
+
+    with pytest.raises(ValueError) as exc_info:
+        llm.invoke(messages)
+    assert "Unknown role: invalid_role" in str(exc_info.value)
+
+
+def test_cohere_llm_get_messages_v2_all_roles(mock_cohere: Mock) -> None:
+    """Test get_messages_v2 method handles all message roles correctly."""
+    # Mock Cohere message types
+    mock_system_msg = MagicMock()
+    mock_user_msg = MagicMock()
+    mock_assistant_msg = MagicMock()
+
+    mock_cohere.SystemChatMessageV2.return_value = mock_system_msg
+    mock_cohere.UserChatMessageV2.return_value = mock_user_msg
+    mock_cohere.AssistantChatMessageV2.return_value = mock_assistant_msg
+
+    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 = CohereLLM(model_name="something")
+    result_messages = llm.get_messages_v2(messages)
+
+    # Verify the correct number of messages are returned
+    assert len(result_messages) == 4
+
+    # Verify the correct Cohere message constructors were called
+    mock_cohere.SystemChatMessageV2.assert_called_once_with(
+        content="You are a helpful assistant."
+    )
+    assert mock_cohere.UserChatMessageV2.call_count == 2
+    mock_cohere.AssistantChatMessageV2.assert_called_once_with(content="Hi there!")
+
+
+def test_cohere_invoke_v2_with_response_format_raises_error(mock_cohere: 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 = CohereLLM(api_key="test")
+
+    with pytest.raises(NotImplementedError) as exc_info:
+        llm.invoke(messages, response_format=TestModel)
+
+    assert "CohereLLM does not currently support structured output" in str(
+        exc_info.value
+    )
+
+
+def test_cohere_llm_close(mock_cohere: Mock) -> None:
+    llm = CohereLLM(model_name="something")
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        llm.close()
+
+
+@pytest.mark.asyncio
+async def test_cohere_llm_aclose(mock_cohere: Mock) -> None:
+    llm = CohereLLM(model_name="something")
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        await llm.aclose()

참고/neo4j-graphrag-python-main/tests/unit/llm/test_llm_utils.py

@@ -0,0 +1,65 @@
+#  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 pytest
+
+from neo4j_graphrag.llm.utils import (
+    legacy_inputs_to_messages,
+    system_instruction_from_messages,
+)
+from neo4j_graphrag.message_history import InMemoryMessageHistory
+from neo4j_graphrag.types import LLMMessage
+
+
+def test_system_instruction_from_messages_found() -> None:
+    messages = [
+        LLMMessage(role="system", content="You are helpful"),
+        LLMMessage(role="user", content="hi"),
+    ]
+    assert system_instruction_from_messages(messages) == "You are helpful"
+
+
+def test_system_instruction_from_messages_not_found() -> None:
+    messages = [LLMMessage(role="user", content="hi")]
+    assert system_instruction_from_messages(messages) is None
+
+
+def test_legacy_inputs_with_message_history_instance() -> None:
+    history = InMemoryMessageHistory()
+    history.add_message(LLMMessage(role="user", content="previous"))
+    result = legacy_inputs_to_messages("follow-up", message_history=history)
+    assert result[-1]["content"] == "follow-up"
+    assert result[0]["content"] == "previous"
+
+
+def test_legacy_inputs_system_instruction_conflict_warns() -> None:
+    messages = [LLMMessage(role="system", content="existing")]
+    with pytest.warns(UserWarning, match="system_instruction provided but ignored"):
+        legacy_inputs_to_messages(
+            "hi", message_history=messages, system_instruction="new"
+        )
+
+
+def test_legacy_inputs_prompt_as_list() -> None:
+    prompt_list = [LLMMessage(role="user", content="hello")]
+    result = legacy_inputs_to_messages(prompt_list)
+    assert result == prompt_list
+
+
+def test_legacy_inputs_prompt_as_message_history() -> None:
+    history = InMemoryMessageHistory()
+    history.add_message(LLMMessage(role="user", content="from history"))
+    result = legacy_inputs_to_messages(history)
+    assert len(result) == 1
+    assert result[0]["content"] == "from history"

참고/neo4j-graphrag-python-main/tests/unit/llm/test_mistralai_llm.py

@@ -0,0 +1,518 @@
+#  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()

참고/neo4j-graphrag-python-main/tests/unit/llm/test_ollama_llm.py

@@ -0,0 +1,700 @@
+#  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, cast
+from unittest.mock import MagicMock, Mock, patch
+
+import ollama
+import pytest
+from neo4j_graphrag.exceptions import LLMGenerationError
+from neo4j_graphrag.llm import LLMResponse
+from neo4j_graphrag.llm.ollama_llm import OllamaLLM
+from neo4j_graphrag.llm.types import ToolCallResponse
+from neo4j_graphrag.tool import Tool
+from neo4j_graphrag.types import LLMMessage
+from pydantic import BaseModel, ConfigDict
+
+
+def get_mock_ollama() -> MagicMock:
+    mock = MagicMock()
+    mock.ResponseError = ollama.ResponseError
+    return mock
+
+
+def _as_mock(value: Any) -> MagicMock:
+    return cast(MagicMock, value)
+
+
+@patch("builtins.__import__", side_effect=ImportError)
+def test_ollama_llm_missing_dependency(mock_import: Mock) -> None:
+    with pytest.raises(ImportError):
+        OllamaLLM(model_name="llama3.2")
+
+
+@patch("builtins.__import__")
+def test_ollama_llm_happy_path_deprecated_options(mock_import: Mock) -> None:
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+    mock_ollama.Client.return_value.chat.return_value = MagicMock(
+        message=MagicMock(content="ollama chat response"),
+    )
+    model = "gpt"
+    model_params = {"temperature": 0.3}
+    with pytest.warns(DeprecationWarning) as record:
+        llm = OllamaLLM(
+            model,
+            model_params=model_params,
+        )
+    assert len(record) == 1
+    assert isinstance(record[0].message, Warning)
+    assert (
+        'you must use model_params={"options": {"temperature": 0}}'
+        in record[0].message.args[0]
+    )
+
+    question = "What is graph RAG?"
+    res = llm.invoke(question)
+    assert isinstance(res, LLMResponse)
+    assert res.content == "ollama chat response"
+    messages = [
+        {"role": "user", "content": question},
+    ]
+    _as_mock(llm.client.chat).assert_called_once_with(
+        model=model, messages=messages, options={"temperature": 0.3}
+    )
+
+
+@patch("builtins.__import__")
+def test_ollama_llm_unsupported_streaming(mock_import: Mock) -> None:
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+    mock_ollama.Client.return_value.chat.return_value = MagicMock(
+        message=MagicMock(content="ollama chat response"),
+    )
+    model = "gpt"
+    model_params = {"stream": True}
+    with pytest.raises(ValueError):
+        OllamaLLM(
+            model,
+            model_params=model_params,
+        )
+
+
+@patch("builtins.__import__")
+def test_ollama_llm_happy_path(mock_import: Mock) -> None:
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+    mock_ollama.Client.return_value.chat.return_value = MagicMock(
+        message=MagicMock(content="ollama chat response"),
+    )
+    model = "gpt"
+    options = {"temperature": 0.3}
+    model_params = {"options": options, "format": "json"}
+    question = "What is graph RAG?"
+    llm = OllamaLLM(
+        model_name=model,
+        model_params=model_params,
+    )
+    res = llm.invoke(question)
+    assert isinstance(res, LLMResponse)
+    assert res.content == "ollama chat response"
+    messages = [
+        {"role": "user", "content": question},
+    ]
+    _as_mock(llm.client.chat).assert_called_once_with(
+        model=model,
+        messages=messages,
+        options=options,
+        format="json",
+    )
+
+
+@patch("builtins.__import__")
+def test_ollama_invoke_with_system_instruction_happy_path(mock_import: Mock) -> None:
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+    mock_ollama.Client.return_value.chat.return_value = MagicMock(
+        message=MagicMock(content="ollama chat response"),
+    )
+    model = "gpt"
+    options = {"temperature": 0.3}
+    model_params = {"options": options, "format": "json"}
+    llm = OllamaLLM(
+        model,
+        model_params=model_params,
+    )
+    system_instruction = "You are a helpful assistant."
+    question = "What about next season?"
+
+    response = llm.invoke(question, system_instruction=system_instruction)
+    assert response.content == "ollama chat response"
+    messages = [{"role": "system", "content": system_instruction}]
+    messages.append({"role": "user", "content": question})
+    _as_mock(llm.client.chat).assert_called_once_with(
+        model=model,
+        messages=messages,
+        options=options,
+        format="json",
+    )
+
+
+@patch("builtins.__import__")
+def test_ollama_invoke_with_message_history_happy_path(mock_import: Mock) -> None:
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+    mock_ollama.Client.return_value.chat.return_value = MagicMock(
+        message=MagicMock(content="ollama chat response"),
+    )
+    model = "gpt"
+    options = {"temperature": 0.3}
+    model_params = {"options": options}
+    llm = OllamaLLM(
+        model,
+        model_params=model_params,
+    )
+    message_history = [
+        {"role": "user", "content": "When does the sun come up in the summer?"},
+        {"role": "assistant", "content": "Usually around 6am."},
+    ]
+    question = "What about next season?"
+
+    response = llm.invoke(question, message_history)  # type: ignore
+    assert response.content == "ollama chat response"
+    messages = [m for m in message_history]
+    messages.append({"role": "user", "content": question})
+    _as_mock(llm.client.chat).assert_called_once_with(
+        model=model, messages=messages, options=options
+    )
+
+
+@patch("builtins.__import__")
+def test_ollama_invoke_with_message_history_and_system_instruction(
+    mock_import: Mock,
+) -> None:
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+    mock_ollama.Client.return_value.chat.return_value = MagicMock(
+        message=MagicMock(content="ollama chat response"),
+    )
+    model = "gpt"
+    options = {"temperature": 0.3}
+    model_params = {"options": options}
+    system_instruction = "You are a helpful assistant."
+    llm = OllamaLLM(
+        model,
+        model_params=model_params,
+    )
+    message_history = [
+        {"role": "user", "content": "When does the sun come up in the summer?"},
+        {"role": "assistant", "content": "Usually around 6am."},
+    ]
+    question = "What about next season?"
+
+    response = llm.invoke(
+        question,
+        message_history,  # type: ignore
+        system_instruction=system_instruction,
+    )
+    assert response.content == "ollama chat response"
+    messages = [{"role": "system", "content": system_instruction}]
+    messages.extend(message_history)
+    messages.append({"role": "user", "content": question})
+    _as_mock(llm.client.chat).assert_called_once_with(
+        model=model, messages=messages, options=options
+    )
+    assert _as_mock(llm.client.chat).call_count == 1
+
+
+@patch("builtins.__import__")
+def test_ollama_invoke_with_message_history_validation_error(mock_import: Mock) -> None:
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+    mock_ollama.ResponseError = ollama.ResponseError
+    model = "gpt"
+    options = {"temperature": 0.3}
+    model_params = {"options": options}
+    system_instruction = "You are a helpful assistant."
+    llm = OllamaLLM(
+        model,
+        model_params=model_params,
+        system_instruction=system_instruction,
+    )
+    message_history = [
+        {"role": "human", "content": "When does the sun come up in the summer?"},
+        {"role": "assistant", "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("builtins.__import__")
+async def test_ollama_ainvoke_happy_path(mock_import: Mock) -> None:
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+
+    async def mock_chat_async(*_args: Any, **_kwargs: Any) -> MagicMock:
+        return MagicMock(
+            message=MagicMock(content="ollama chat response"),
+        )
+
+    mock_ollama.AsyncClient.return_value.chat = mock_chat_async
+    model = "gpt"
+    options = {"temperature": 0.3}
+    model_params = {"options": options}
+    question = "What is graph RAG?"
+    llm = OllamaLLM(
+        model,
+        model_params=model_params,
+    )
+
+    res = await llm.ainvoke(question)
+    assert isinstance(res, LLMResponse)
+    assert res.content == "ollama chat response"
+
+
+# V2 Interface Tests
+@patch("builtins.__import__")
+def test_ollama_llm_invoke_v2_happy_path(mock_import: Mock) -> None:
+    """Test V2 interface invoke method with List[LLMMessage] input."""
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+    mock_ollama.Client.return_value.chat.return_value = MagicMock(
+        message=MagicMock(content="ollama v2 response"),
+    )
+    mock_ollama.Message = MagicMock()
+
+    model = "llama2"
+    options = {"temperature": 0.3}
+    model_params = {"options": options}
+
+    messages: list[LLMMessage] = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "What is graph RAG?"},
+    ]
+
+    llm = OllamaLLM(
+        model_name=model,
+        model_params=model_params,
+    )
+    res = llm.invoke(messages)
+
+    assert isinstance(res, LLMResponse)
+    assert res.content == "ollama v2 response"
+
+    # Verify get_brand_new_messages was called correctly
+    assert mock_ollama.Message.call_count == 2
+    mock_ollama.Message.assert_any_call(**messages[0])
+    mock_ollama.Message.assert_any_call(**messages[1])
+
+    # Verify the client was called with correct parameters
+    _as_mock(llm.client.chat).assert_called_once_with(
+        model=model,
+        messages=[mock_ollama.Message.return_value, mock_ollama.Message.return_value],
+        options=options,
+    )
+
+
+@pytest.mark.asyncio
+@patch("builtins.__import__")
+async def test_ollama_llm_ainvoke_v2_happy_path(mock_import: Mock) -> None:
+    """Test V2 interface ainvoke method with List[LLMMessage] input."""
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+    mock_ollama.Message = MagicMock()
+
+    async def mock_chat_async(*_args: Any, **_kwargs: Any) -> MagicMock:
+        return MagicMock(
+            message=MagicMock(content="ollama async v2 response"),
+        )
+
+    mock_ollama.AsyncClient.return_value.chat = mock_chat_async
+
+    model = "llama2"
+    options = {"temperature": 0.5}
+    model_params = {"options": options}
+
+    messages: list[LLMMessage] = [
+        {"role": "user", "content": "What is Neo4j?"},
+        {"role": "assistant", "content": "Neo4j is a graph database."},
+        {"role": "user", "content": "How does it work?"},
+    ]
+
+    llm = OllamaLLM(
+        model_name=model,
+        model_params=model_params,
+    )
+    res = await llm.ainvoke(messages)
+
+    assert isinstance(res, LLMResponse)
+    assert res.content == "ollama async v2 response"
+
+    # Verify get_brand_new_messages was called correctly
+    assert mock_ollama.Message.call_count == 3
+    for message in messages:
+        mock_ollama.Message.assert_any_call(**message)
+
+
+@patch("builtins.__import__")
+def test_ollama_llm_invoke_v2_error_handling(mock_import: Mock) -> None:
+    """Test V2 interface error handling when OllamaResponseError occurs."""
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+    mock_ollama.Client.return_value.chat.side_effect = ollama.ResponseError(
+        "Ollama error"
+    )
+    mock_ollama.Message = MagicMock()
+
+    model = "llama2"
+    messages: list[LLMMessage] = [
+        {"role": "user", "content": "This will cause an error."},
+    ]
+
+    llm = OllamaLLM(model_name=model)
+
+    with pytest.raises(LLMGenerationError):
+        llm.invoke(messages)
+
+
+@pytest.mark.asyncio
+@patch("builtins.__import__")
+async def test_ollama_llm_ainvoke_v2_error_handling(mock_import: Mock) -> None:
+    """Test V2 interface async error handling when OllamaResponseError occurs."""
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+    mock_ollama.Message = MagicMock()
+
+    async def mock_chat_async_error(*_args: Any, **_kwargs: Any) -> None:
+        raise ollama.ResponseError("Async Ollama error")
+
+    mock_ollama.AsyncClient.return_value.chat = mock_chat_async_error
+
+    model = "llama2"
+    messages: list[LLMMessage] = [
+        {"role": "user", "content": "This will cause an async error."},
+    ]
+
+    llm = OllamaLLM(model_name=model)
+
+    with pytest.raises(LLMGenerationError):
+        await llm.ainvoke(messages)
+
+
+@patch("builtins.__import__")
+def test_ollama_llm_input_type_switching_string(mock_import: Mock) -> None:
+    """Test that string input correctly routes to legacy invoke method."""
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+    mock_ollama.Client.return_value.chat.return_value = MagicMock(
+        message=MagicMock(content="legacy response"),
+    )
+
+    model = "llama2"
+    question = "What is graph RAG?"
+
+    llm = OllamaLLM(model_name=model)
+    res = llm.invoke(question)
+
+    assert isinstance(res, LLMResponse)
+    assert res.content == "legacy response"
+
+    # Verify legacy method was used (messages should be built via get_messages)
+    _as_mock(llm.client.chat).assert_called_once()
+    call_args = _as_mock(llm.client.chat).call_args[1]
+    assert call_args["model"] == model
+    assert len(call_args["messages"]) == 1
+    assert call_args["messages"][0]["role"] == "user"
+    assert call_args["messages"][0]["content"] == question
+
+
+@patch("builtins.__import__")
+def test_ollama_llm_input_type_switching_list(mock_import: Mock) -> None:
+    """Test that List[LLMMessage] input correctly routes to V2 invoke method."""
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+    mock_ollama.Client.return_value.chat.return_value = MagicMock(
+        message=MagicMock(content="v2 response"),
+    )
+    mock_ollama.Message = MagicMock()
+
+    model = "llama2"
+    messages: list[LLMMessage] = [
+        {"role": "user", "content": "What is graph RAG?"},
+    ]
+
+    llm = OllamaLLM(model_name=model)
+    res = llm.invoke(messages)
+
+    assert isinstance(res, LLMResponse)
+    assert res.content == "v2 response"
+
+    # Verify V2 method was used (ollama.Message should be called)
+    mock_ollama.Message.assert_called_once_with(**messages[0])
+
+
+@patch("builtins.__import__")
+def test_ollama_llm_invalid_input_type(mock_import: Mock) -> None:
+    """Test that invalid input type raises ValueError."""
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+
+    llm = OllamaLLM(model_name="llama2")
+
+    # Test with invalid input type (neither string nor list)
+    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("builtins.__import__")
+async def test_ollama_llm_ainvoke_invalid_input_type(mock_import: Mock) -> None:
+    """Test that invalid input type raises ValueError in async method."""
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+
+    llm = OllamaLLM(model_name="llama2")
+
+    # Test with invalid input type (neither string nor list)
+    with pytest.raises(ValueError) as exc_info:
+        await llm.ainvoke({"invalid": "dict"})  # type: ignore
+    assert "Invalid input type for ainvoke method" in str(exc_info.value)
+
+
+@patch("builtins.__import__")
+def test_ollama_llm_get_brand_new_messages_all_roles(mock_import: Mock) -> None:
+    """Test get_brand_new_messages method handles all message roles correctly."""
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+    mock_ollama.Message = MagicMock()
+
+    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 = OllamaLLM(model_name="llama2")
+    result_messages = llm.get_messages_v2(messages)
+
+    # Convert to list for easier testing
+    result_list = list(result_messages)
+
+    # Verify correct number of ollama.Message objects created
+    assert len(result_list) == 4
+    assert mock_ollama.Message.call_count == 4
+
+    # Verify each message was converted properly
+    for message in messages:
+        mock_ollama.Message.assert_any_call(**message)
+
+
+@patch("builtins.__import__")
+def test_ollama_llm_invoke_with_tools_happy_path(
+    mock_import: Mock,
+    test_tool: Tool,
+) -> None:
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+
+    # Mock the tool call response
+    mock_function = MagicMock()
+    mock_function.name = "test_tool"
+    mock_function.arguments = {"param1": "value1"}
+
+    mock_tool_call = MagicMock()
+    mock_tool_call.function = mock_function
+
+    mock_ollama.Client.return_value.chat.return_value = MagicMock(
+        message=MagicMock(content="ollama tool response", tool_calls=[mock_tool_call])
+    )
+
+    llm = OllamaLLM(model_name="gpt", model_params={"options": {"temperature": 0}})
+    tools = [test_tool]
+
+    res = llm.invoke_with_tools("my text", tools)
+    assert isinstance(res, ToolCallResponse)
+    assert len(res.tool_calls) == 1
+    assert res.tool_calls[0].name == "test_tool"
+    assert res.tool_calls[0].arguments == {"param1": "value1"}
+    assert res.content == "ollama tool response"
+
+
+@patch("builtins.__import__")
+def test_ollama_llm_invoke_with_tools_with_message_history(
+    mock_import: Mock,
+    test_tool: Tool,
+) -> None:
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+
+    # Mock the tool call response
+    mock_function = MagicMock()
+    mock_function.name = "test_tool"
+    mock_function.arguments = {"param1": "value1"}
+
+    mock_tool_call = MagicMock()
+    mock_tool_call.function = mock_function
+    mock_ollama.Client.return_value.chat.return_value = MagicMock(
+        message=MagicMock(content="ollama tool response", tool_calls=[mock_tool_call])
+    )
+    llm = OllamaLLM(
+        api_key="my key", model_name="gpt", model_params={"options": {"temperature": 0}}
+    )
+    tools = [test_tool]
+
+    message_history = [
+        {"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_with_tools(question, tools, message_history)  # type: ignore
+    assert isinstance(res, ToolCallResponse)
+    assert len(res.tool_calls) == 1
+    assert res.tool_calls[0].name == "test_tool"
+    assert res.tool_calls[0].arguments == {"param1": "value1"}
+
+    # Verify the correct messages were passed
+    message_history.append({"role": "user", "content": question})
+    # Use assert_called_once() instead of assert_called_once_with() to avoid issues with overloaded functions
+    _as_mock(llm.client.chat).assert_called_once()
+    # Check call arguments individually
+    call_args = _as_mock(llm.client.chat).call_args[1]  # Get the keyword arguments
+    assert call_args["messages"] == message_history
+    assert call_args["model"] == "gpt"
+    # Check tools content rather than direct equality
+    assert len(call_args["tools"]) == 1
+    assert call_args["tools"][0]["type"] == "function"
+    assert call_args["tools"][0]["function"]["name"] == "test_tool"
+    assert call_args["tools"][0]["function"]["description"] == "A test tool"
+
+
+@patch("builtins.__import__")
+def test_ollama_llm_invoke_with_tools_with_system_instruction(
+    mock_import: Mock,
+    test_tool: Mock,
+) -> None:
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+
+    # Mock the tool call response
+    mock_function = MagicMock()
+    mock_function.name = "test_tool"
+    mock_function.arguments = {"param1": "value1"}
+
+    mock_tool_call = MagicMock()
+    mock_tool_call.function = mock_function
+
+    mock_ollama.Client.return_value.chat.return_value = MagicMock(
+        message=MagicMock(content="ollama tool response", tool_calls=[mock_tool_call])
+    )
+
+    llm = OllamaLLM(
+        api_key="my key", model_name="gpt", model_params={"options": {"temperature": 0}}
+    )
+    tools = [test_tool]
+
+    system_instruction = "You are a helpful assistant."
+
+    res = llm.invoke_with_tools("my text", tools, system_instruction=system_instruction)
+    assert isinstance(res, ToolCallResponse)
+
+    # Verify system instruction was included
+    messages = [{"role": "system", "content": system_instruction}]
+    messages.append({"role": "user", "content": "my text"})
+    # Use assert_called_once() instead of assert_called_once_with() to avoid issues with overloaded functions
+    _as_mock(llm.client.chat).assert_called_once()
+    # Check call arguments individually
+    call_args = _as_mock(llm.client.chat).call_args[1]  # Get the keyword arguments
+    assert call_args["messages"] == messages
+    assert call_args["model"] == "gpt"
+    # Check tools content rather than direct equality
+    assert len(call_args["tools"]) == 1
+    assert call_args["tools"][0]["type"] == "function"
+    assert call_args["tools"][0]["function"]["name"] == "test_tool"
+    assert call_args["tools"][0]["function"]["description"] == "A test tool"
+
+
+@patch("builtins.__import__")
+def test_ollama_llm_invoke_with_tools_error(mock_import: Mock, test_tool: Tool) -> None:
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+
+    # Mock an Ollama response error
+    mock_ollama.Client.return_value.chat.side_effect = ollama.ResponseError(
+        "Test error"
+    )
+
+    llm = OllamaLLM(
+        api_key="my key", model_name="gpt", model_params={"options": {"temperature": 0}}
+    )
+    tools = [test_tool]
+
+    with pytest.raises(LLMGenerationError):
+        llm.invoke_with_tools("my text", tools)
+
+
+class _TestModelForOllama(BaseModel):
+    """Test model for structured output tests."""
+
+    model_config = ConfigDict(extra="forbid")
+    value: str
+
+
+@patch("builtins.__import__")
+def test_ollama_invoke_v2_with_response_format_raises_error(mock_import: Mock) -> None:
+    """Test V2 interface raises NotImplementedError when response_format is used."""
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+
+    messages: List[LLMMessage] = [{"role": "user", "content": "Test"}]
+    llm = OllamaLLM(model_name="llama2")
+
+    with pytest.raises(NotImplementedError) as exc_info:
+        llm.invoke(messages, response_format=_TestModelForOllama)
+
+    assert "OllamaLLM does not currently support structured output" in str(
+        exc_info.value
+    )
+
+
+@patch("builtins.__import__")
+def test_ollama_llm_close(mock_import: Mock) -> None:
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+
+    llm = OllamaLLM(model_name="llama3.2", model_params={"options": {}})
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        llm.close()
+
+
+@pytest.mark.asyncio
+@patch("builtins.__import__")
+async def test_ollama_llm_aclose(mock_import: Mock) -> None:
+    mock_ollama = get_mock_ollama()
+    mock_import.return_value = mock_ollama
+
+    llm = OllamaLLM(model_name="llama3.2", model_params={"options": {}})
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        await llm.aclose()

참고/neo4j-graphrag-python-main/tests/unit/llm/test_openai_llm.py

@@ -0,0 +1,977 @@
+#  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 unittest.mock import AsyncMock, MagicMock, Mock, patch
+from typing import Any, Callable, List
+import builtins
+
+import httpx
+import openai
+import pytest
+from neo4j_graphrag.exceptions import LLMGenerationError
+from neo4j_graphrag.llm.types import LLMResponse
+from neo4j_graphrag.llm.openai_llm import AzureOpenAILLM, OpenAILLM
+from neo4j_graphrag.llm.types import ToolCallResponse
+from neo4j_graphrag.tool import Tool
+from neo4j_graphrag.types import LLMMessage
+from pydantic import BaseModel, ConfigDict
+
+# Save the original __import__ before any patches are applied
+_original_import = builtins.__import__
+
+
+def get_mock_openai() -> MagicMock:
+    mock = MagicMock()
+    mock.OpenAIError = openai.OpenAIError
+    mock.httpx = httpx
+    return mock
+
+
+def create_selective_import_mock(mock_openai: MagicMock) -> Callable[..., Any]:
+    """Create a mock that only intercepts 'openai' imports, letting others pass through."""
+
+    def selective_import(name: str, *args: Any, **kwargs: Any) -> Any:
+        if name == "openai":
+            return mock_openai
+        return _original_import(name, *args, **kwargs)
+
+    return selective_import
+
+
+@patch("builtins.__import__", side_effect=ImportError)
+def test_openai_llm_missing_dependency(_mock_import: Mock) -> None:
+    with pytest.raises(ImportError):
+        OpenAILLM(model_name="gpt-5")
+
+
+@patch("builtins.__import__")
+def test_openai_llm_happy_path(mock_import: Mock) -> None:
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+    mock_openai.OpenAI.return_value.chat.completions.create.return_value = MagicMock(
+        choices=[MagicMock(message=MagicMock(content="openai chat response"))],
+    )
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+
+    res = llm.invoke("my text")
+    assert isinstance(res, LLMResponse)
+    assert res.content == "openai chat response"
+
+
+@patch("builtins.__import__")
+def test_openai_llm_with_message_history_happy_path(mock_import: Mock) -> None:
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+    mock_openai.OpenAI.return_value.chat.completions.create.return_value = MagicMock(
+        choices=[MagicMock(message=MagicMock(content="openai chat response"))],
+    )
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+    message_history = [
+        {"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)  # type: ignore
+    assert isinstance(res, LLMResponse)
+    assert res.content == "openai chat response"
+    message_history.append({"role": "user", "content": question})
+    # Use assert_called_once() instead of assert_called_once_with() to avoid issues with overloaded functions
+    llm.client.chat.completions.create.assert_called_once()  # type: ignore
+    # Check call arguments individually
+    call_args = llm.client.chat.completions.create.call_args[  # type: ignore
+        1
+    ]  # Get the keyword arguments
+    assert call_args["messages"] == message_history
+    assert call_args["model"] == "gpt"
+
+
+@patch("builtins.__import__")
+def test_openai_llm_with_message_history_and_system_instruction(
+    mock_import: Mock,
+) -> None:
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+    mock_openai.OpenAI.return_value.chat.completions.create.return_value = MagicMock(
+        choices=[MagicMock(message=MagicMock(content="openai chat response"))],
+    )
+    system_instruction = "You are a helpful assistent."
+    llm = OpenAILLM(
+        api_key="my key",
+        model_name="gpt",
+    )
+    message_history = [
+        {"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)  # type: ignore
+    assert isinstance(res, LLMResponse)
+    assert res.content == "openai chat response"
+    messages = [{"role": "system", "content": system_instruction}]
+    messages.extend(message_history)
+    messages.append({"role": "user", "content": question})
+    # Use assert_called_once() instead of assert_called_once_with() to avoid issues with overloaded functions
+    llm.client.chat.completions.create.assert_called_once()  # type: ignore
+    # Check call arguments individually
+    call_args = llm.client.chat.completions.create.call_args[  # type: ignore
+        1
+    ]  # Get the keyword arguments
+    assert call_args["messages"] == messages
+    assert call_args["model"] == "gpt"
+
+    assert llm.client.chat.completions.create.call_count == 1  # type: ignore
+
+
+@patch("builtins.__import__")
+def test_openai_llm_with_message_history_validation_error(mock_import: Mock) -> None:
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+    mock_openai.OpenAI.return_value.chat.completions.create.return_value = MagicMock(
+        choices=[MagicMock(message=MagicMock(content="openai chat response"))],
+    )
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+    message_history = [
+        {"role": "human", "content": "When does the sun come up in the summer?"},
+        {"role": "assistant", "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)
+
+
+@patch("builtins.__import__")
+@patch("json.loads")
+def test_openai_llm_invoke_with_tools_happy_path(
+    mock_json_loads: Mock,
+    mock_import: Mock,
+    test_tool: Tool,
+) -> None:
+    # Set up json.loads to return a dictionary
+    mock_json_loads.return_value = {"param1": "value1"}
+
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    # Mock the tool call response
+    mock_function = MagicMock()
+    mock_function.name = "test_tool"
+    mock_function.arguments = '{"param1": "value1"}'
+
+    mock_tool_call = MagicMock()
+    mock_tool_call.function = mock_function
+
+    mock_openai.OpenAI.return_value.chat.completions.create.return_value = MagicMock(
+        choices=[
+            MagicMock(
+                message=MagicMock(
+                    content="openai tool response", tool_calls=[mock_tool_call]
+                )
+            )
+        ],
+    )
+
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+    tools = [test_tool]
+
+    res = llm.invoke_with_tools("my text", tools)
+    assert isinstance(res, ToolCallResponse)
+    assert len(res.tool_calls) == 1
+    assert res.tool_calls[0].name == "test_tool"
+    assert res.tool_calls[0].arguments == {"param1": "value1"}
+    assert res.content == "openai tool response"
+
+
+@patch("builtins.__import__")
+@patch("json.loads")
+def test_openai_llm_invoke_with_tools_with_message_history(
+    mock_json_loads: Mock,
+    mock_import: Mock,
+    test_tool: Tool,
+) -> None:
+    # Set up json.loads to return a dictionary
+    mock_json_loads.return_value = {"param1": "value1"}
+
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    # Mock the tool call response
+    mock_function = MagicMock()
+    mock_function.name = "test_tool"
+    mock_function.arguments = '{"param1": "value1"}'
+
+    mock_tool_call = MagicMock()
+    mock_tool_call.function = mock_function
+
+    mock_openai.OpenAI.return_value.chat.completions.create.return_value = MagicMock(
+        choices=[
+            MagicMock(
+                message=MagicMock(
+                    content="openai tool response", tool_calls=[mock_tool_call]
+                )
+            )
+        ],
+    )
+
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+    tools = [test_tool]
+
+    message_history = [
+        {"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_with_tools(question, tools, message_history)  # type: ignore
+    assert isinstance(res, ToolCallResponse)
+    assert len(res.tool_calls) == 1
+    assert res.tool_calls[0].name == "test_tool"
+    assert res.tool_calls[0].arguments == {"param1": "value1"}
+
+    # Verify the correct messages were passed
+    message_history.append({"role": "user", "content": question})
+    # Use assert_called_once() instead of assert_called_once_with() to avoid issues with overloaded functions
+    llm.client.chat.completions.create.assert_called_once()  # type: ignore
+    # Check call arguments individually
+    call_args = llm.client.chat.completions.create.call_args[  # type: ignore
+        1
+    ]  # Get the keyword arguments
+    assert call_args["messages"] == message_history
+    assert call_args["model"] == "gpt"
+    # Check tools content rather than direct equality
+    assert len(call_args["tools"]) == 1
+    assert call_args["tools"][0]["type"] == "function"
+    assert call_args["tools"][0]["function"]["name"] == "test_tool"
+    assert call_args["tools"][0]["function"]["description"] == "A test tool"
+    assert call_args["tool_choice"] == "auto"
+    assert call_args["temperature"] == 0.0
+
+
+@patch("builtins.__import__")
+@patch("json.loads")
+def test_openai_llm_invoke_with_tools_with_system_instruction(
+    mock_json_loads: Mock,
+    mock_import: Mock,
+    test_tool: Mock,
+) -> None:
+    # Set up json.loads to return a dictionary
+    mock_json_loads.return_value = {"param1": "value1"}
+
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    # Mock the tool call response
+    mock_function = MagicMock()
+    mock_function.name = "test_tool"
+    mock_function.arguments = '{"param1": "value1"}'
+
+    mock_tool_call = MagicMock()
+    mock_tool_call.function = mock_function
+
+    mock_openai.OpenAI.return_value.chat.completions.create.return_value = MagicMock(
+        choices=[
+            MagicMock(
+                message=MagicMock(
+                    content="openai tool response", tool_calls=[mock_tool_call]
+                )
+            )
+        ],
+    )
+
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+    tools = [test_tool]
+
+    system_instruction = "You are a helpful assistant."
+
+    res = llm.invoke_with_tools("my text", tools, system_instruction=system_instruction)
+    assert isinstance(res, ToolCallResponse)
+
+    # Verify system instruction was included
+    messages = [{"role": "system", "content": system_instruction}]
+    messages.append({"role": "user", "content": "my text"})
+    # Use assert_called_once() instead of assert_called_once_with() to avoid issues with overloaded functions
+    llm.client.chat.completions.create.assert_called_once()  # type: ignore
+    # Check call arguments individually
+    call_args = llm.client.chat.completions.create.call_args[  # type: ignore
+        1
+    ]  # Get the keyword arguments
+    assert call_args["messages"] == messages
+    assert call_args["model"] == "gpt"
+    # Check tools content rather than direct equality
+    assert len(call_args["tools"]) == 1
+    assert call_args["tools"][0]["type"] == "function"
+    assert call_args["tools"][0]["function"]["name"] == "test_tool"
+    assert call_args["tools"][0]["function"]["description"] == "A test tool"
+    assert call_args["tool_choice"] == "auto"
+    assert call_args["temperature"] == 0.0
+
+
+@patch("builtins.__import__")
+def test_openai_llm_invoke_with_tools_error(mock_import: Mock, test_tool: Tool) -> None:
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    # Mock an OpenAI error
+    mock_openai.OpenAI.return_value.chat.completions.create.side_effect = (
+        openai.OpenAIError("Test error")
+    )
+
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+    tools = [test_tool]
+
+    with pytest.raises(LLMGenerationError):
+        llm.invoke_with_tools("my text", tools)
+
+
+@patch("builtins.__import__", side_effect=ImportError)
+def test_azure_openai_llm_missing_dependency(_mock_import: Mock) -> None:
+    with pytest.raises(ImportError):
+        AzureOpenAILLM(model_name="gpt-5")
+
+
+@patch("builtins.__import__")
+def test_azure_openai_llm_happy_path(mock_import: Mock) -> None:
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+    mock_openai.AzureOpenAI.return_value.chat.completions.create.return_value = (
+        MagicMock(
+            choices=[MagicMock(message=MagicMock(content="openai chat response"))],
+        )
+    )
+    llm = AzureOpenAILLM(
+        model_name="gpt",
+        azure_endpoint="https://test.openai.azure.com/",
+        api_key="my key",
+        api_version="version",
+    )
+
+    res = llm.invoke("my text")
+    assert isinstance(res, LLMResponse)
+    assert res.content == "openai chat response"
+
+
+@patch("builtins.__import__")
+def test_azure_openai_llm_with_message_history_happy_path(mock_import: Mock) -> None:
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+    mock_openai.AzureOpenAI.return_value.chat.completions.create.return_value = (
+        MagicMock(
+            choices=[MagicMock(message=MagicMock(content="openai chat response"))],
+        )
+    )
+    llm = AzureOpenAILLM(
+        model_name="gpt",
+        azure_endpoint="https://test.openai.azure.com/",
+        api_key="my key",
+        api_version="version",
+    )
+
+    message_history = [
+        {"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)  # type: ignore
+    assert isinstance(res, LLMResponse)
+    assert res.content == "openai chat response"
+    message_history.append({"role": "user", "content": question})
+    # Use assert_called_once() instead of assert_called_once_with() to avoid issues with overloaded functions
+    llm.client.chat.completions.create.assert_called_once()  # type: ignore
+    # Check call arguments individually
+    call_args = llm.client.chat.completions.create.call_args[  # type: ignore
+        1
+    ]  # Get the keyword arguments
+    assert call_args["messages"] == message_history
+    assert call_args["model"] == "gpt"
+
+
+@patch("builtins.__import__")
+def test_azure_openai_llm_with_message_history_validation_error(
+    mock_import: Mock,
+) -> None:
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+    mock_openai.AzureOpenAI.return_value.chat.completions.create.return_value = (
+        MagicMock(
+            choices=[MagicMock(message=MagicMock(content="openai chat response"))],
+        )
+    )
+    llm = AzureOpenAILLM(
+        model_name="gpt",
+        azure_endpoint="https://test.openai.azure.com/",
+        api_key="my key",
+        api_version="version",
+    )
+
+    message_history = [
+        {"role": "user", "content": 33},
+    ]
+    question = "What about next season?"
+
+    with pytest.raises(LLMGenerationError) as exc_info:
+        llm.invoke(question, message_history)  # type: ignore
+    assert "Input should be a valid string" in str(exc_info.value)
+
+
+@pytest.mark.asyncio
+@patch("builtins.__import__")
+async def test_openai_llm_ainvoke_happy_path(mock_import: Mock) -> None:
+    """Test that ainvoke properly awaits the async call and returns LLMResponse."""
+    # Mock OpenAI module
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    # Build mock response matching OpenAI's structure
+    mock_message = MagicMock()
+    mock_message.content = "Return text"
+
+    mock_choice = MagicMock()
+    mock_choice.message = mock_message
+
+    mock_response = MagicMock()
+    mock_response.choices = [mock_choice]
+
+    # Async function instead of AsyncMock
+    async def async_create(*args, **kwargs):  # type: ignore[no-untyped-def]
+        return mock_response
+
+    mock_openai.AsyncOpenAI.return_value.chat.completions.create = async_create
+
+    model_name = "gpt-3.5-turbo"
+    input_text = "may thy knife chip and shatter"
+    model_params = {"temperature": 0.5}
+    llm = OpenAILLM(model_name, model_params, api_key="test-key")
+
+    response = await llm.ainvoke(input_text)
+
+    # Assert we got the expected content in LLMResponse
+    assert isinstance(response, LLMResponse)
+    assert response.content == "Return text"
+
+
+# LLM Interface V2 Tests
+
+
+@patch("builtins.__import__")
+def test_openai_llm_invoke_v2_happy_path(mock_import: Mock) -> None:
+    """Test V2 interface invoke method with List[LLMMessage] input."""
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+    mock_openai.OpenAI.return_value.chat.completions.create.return_value = MagicMock(
+        choices=[
+            MagicMock(message=MagicMock(content="Paris is the capital of France."))
+        ],
+    )
+
+    messages: List[LLMMessage] = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "What is the capital of France?"},
+    ]
+
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+    response = llm.invoke(messages)
+
+    assert isinstance(response, LLMResponse)
+    assert response.content == "Paris is the capital of France."
+
+    # Verify the client was called correctly
+    llm.client.chat.completions.create.assert_called_once()  # type: ignore
+    call_args = llm.client.chat.completions.create.call_args[1]  # type: ignore
+    # Verify we have the right number of messages and model
+    assert len(call_args["messages"]) == 2
+    assert call_args["model"] == "gpt"
+
+
+@patch("builtins.__import__")
+def test_openai_llm_invoke_v2_with_conversation_history(mock_import: Mock) -> None:
+    """Test V2 interface invoke with conversation history."""
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+    mock_openai.OpenAI.return_value.chat.completions.create.return_value = MagicMock(
+        choices=[
+            MagicMock(message=MagicMock(content="Berlin is the capital of Germany."))
+        ],
+    )
+
+    messages: List[LLMMessage] = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "What is the capital of France?"},
+        {"role": "assistant", "content": "Paris is the capital of France."},
+        {"role": "user", "content": "What about Germany?"},
+    ]
+
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+    response = llm.invoke(messages)
+
+    assert isinstance(response, LLMResponse)
+    assert response.content == "Berlin is the capital of Germany."
+
+    # Verify all messages were passed correctly
+    llm.client.chat.completions.create.assert_called_once()  # type: ignore
+    call_args = llm.client.chat.completions.create.call_args[1]  # type: ignore
+    assert len(call_args["messages"]) == 4
+    assert call_args["model"] == "gpt"
+
+
+@patch("builtins.__import__")
+def test_openai_llm_invoke_v2_no_system_message(mock_import: Mock) -> None:
+    """Test V2 interface invoke without system message."""
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+    mock_openai.OpenAI.return_value.chat.completions.create.return_value = MagicMock(
+        choices=[MagicMock(message=MagicMock(content="I'm doing well, thank you!"))],
+    )
+
+    messages: List[LLMMessage] = [
+        {"role": "user", "content": "Hello, how are you?"},
+    ]
+
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+    response = llm.invoke(messages)
+
+    assert isinstance(response, LLMResponse)
+    assert response.content == "I'm doing well, thank you!"
+
+    # Verify only user message was passed
+    llm.client.chat.completions.create.assert_called_once()  # type: ignore
+    call_args = llm.client.chat.completions.create.call_args[1]  # type: ignore
+    assert len(call_args["messages"]) == 1
+
+
+@pytest.mark.asyncio
+@patch("builtins.__import__")
+async def test_openai_llm_ainvoke_v2_happy_path(mock_import: Mock) -> None:
+    """Test V2 interface async invoke method with List[LLMMessage] input."""
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    # Build mock response matching OpenAI's structure
+    mock_message = MagicMock()
+    mock_message.content = "2+2 equals 4."
+
+    mock_choice = MagicMock()
+    mock_choice.message = mock_message
+
+    mock_response = MagicMock()
+    mock_response.choices = [mock_choice]
+
+    # Async function to simulate .create()
+    async def async_create(*args, **kwargs):  # type: ignore[no-untyped-def]
+        """Async mock for chat completions create."""
+        return mock_response
+
+    mock_openai.AsyncOpenAI.return_value.chat.completions.create = async_create
+
+    messages: List[LLMMessage] = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "What is 2+2?"},
+    ]
+
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+    response = await llm.ainvoke(messages)
+
+    # Assert the returned LLMResponse
+    assert isinstance(response, LLMResponse)
+    assert response.content == "2+2 equals 4."
+
+    # Verify async client was called
+    # Patch async_create itself to track calls
+    called_args = getattr(
+        llm.async_client.chat.completions.create, "__wrapped_args__", None
+    )
+    assert called_args is None or True  # optional, depends on how strict tracking is
+
+
+# Note: Async tool calling test is covered by the synchronous version above
+# The complex mocking of json.loads with local imports makes this test difficult to maintain
+
+
+@patch("builtins.__import__")
+def test_openai_llm_invoke_v2_validation_error(mock_import: Mock) -> None:
+    """Test V2 interface invoke with invalid message format raises error."""
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    messages: List[LLMMessage] = [
+        {"role": "invalid_role", "content": "This should fail."},  # type: ignore
+    ]
+
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+
+    with pytest.raises(ValueError) as exc_info:
+        llm.invoke(messages)
+    assert "Unknown role: invalid_role" in str(exc_info.value)
+
+
+@patch("builtins.__import__")
+def test_openai_llm_get_messages_v2_all_roles(mock_import: Mock) -> None:
+    """Test get_messages_v2 method handles all message roles correctly."""
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    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 = OpenAILLM(api_key="my key", model_name="gpt")
+    result_messages = llm.get_messages_v2(messages)
+
+    # Convert to list for easier testing
+    result_list = list(result_messages)
+
+    # Just verify the correct number of messages are returned
+    # (Detailed content inspection is difficult due to OpenAI message object mocking)
+    assert len(result_list) == 4
+
+
+@patch("builtins.__import__")
+def test_azure_openai_llm_invoke_v2_happy_path(mock_import: Mock) -> None:
+    """Test V2 interface invoke method for Azure OpenAI with List[LLMMessage] input."""
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+    mock_openai.AzureOpenAI.return_value.chat.completions.create.return_value = (
+        MagicMock(
+            choices=[MagicMock(message=MagicMock(content="Azure OpenAI response"))],
+        )
+    )
+
+    messages: List[LLMMessage] = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "What is Azure?"},
+    ]
+
+    llm = AzureOpenAILLM(
+        model_name="gpt",
+        azure_endpoint="https://test.openai.azure.com/",
+        api_key="my key",
+        api_version="version",
+    )
+    response = llm.invoke(messages)
+
+    assert isinstance(response, LLMResponse)
+    assert response.content == "Azure OpenAI response"
+
+    # Verify the correct messages were passed
+    llm.client.chat.completions.create.assert_called_once()  # type: ignore
+    call_args = llm.client.chat.completions.create.call_args[1]  # type: ignore
+    assert len(call_args["messages"]) == 2
+    assert call_args["model"] == "gpt"
+
+
+class _TestModelForOpenAI(BaseModel):
+    """Test model for structured output tests."""
+
+    model_config = ConfigDict(extra="forbid")
+    name: str
+    age: int
+
+
+# JSON schema for structured output tests
+_TEST_JSON_SCHEMA = {
+    "type": "json_schema",
+    "json_schema": {
+        "name": "test_schema",
+        "strict": True,
+        "schema": {"type": "object", "properties": {"result": {"type": "string"}}},
+    },
+}
+
+
+@patch("builtins.__import__")
+def test_openai_llm_invoke_v2_with_pydantic_response_format(mock_import: Mock) -> None:
+    """Test V2 interface with Pydantic model as response_format."""
+
+    mock_openai = get_mock_openai()
+    mock_import.side_effect = create_selective_import_mock(mock_openai)
+    mock_openai.OpenAI.return_value.chat.completions.create.return_value = MagicMock(
+        choices=[MagicMock(message=MagicMock(content='{"name": "John", "age": 30}'))],
+    )
+
+    messages: List[LLMMessage] = [
+        {"role": "user", "content": "Extract person info"},
+    ]
+
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+    response = llm.invoke(messages, response_format=_TestModelForOpenAI)
+
+    assert response.content == '{"name": "John", "age": 30}'
+
+    # Verify the method was called (response_format handling is internal)
+    llm.client.chat.completions.create.assert_called_once()  # type: ignore
+
+
+@patch("builtins.__import__")
+def test_openai_llm_invoke_v2_with_json_schema_response_format(
+    mock_import: Mock,
+) -> None:
+    """Test V2 interface with JSON schema dict as response_format."""
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+    mock_openai.OpenAI.return_value.chat.completions.create.return_value = MagicMock(
+        choices=[MagicMock(message=MagicMock(content='{"result": "success"}'))],
+    )
+
+    messages: List[LLMMessage] = [
+        {"role": "user", "content": "Test"},
+    ]
+
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+    response = llm.invoke(messages, response_format=_TEST_JSON_SCHEMA)
+
+    assert response.content == '{"result": "success"}'
+
+    # Verify the method was called (response_format handling is internal)
+    llm.client.chat.completions.create.assert_called_once()  # type: ignore
+
+
+@pytest.mark.asyncio
+@patch("builtins.__import__")
+async def test_openai_llm_ainvoke_v2_with_pydantic_response_format(
+    mock_import: Mock,
+) -> None:
+    """Test V2 interface async invoke with Pydantic response_format."""
+
+    mock_openai = get_mock_openai()
+    mock_import.side_effect = create_selective_import_mock(mock_openai)
+
+    mock_response = MagicMock()
+    mock_response.choices = [MagicMock(message=MagicMock(content='{"value": "test"}'))]
+
+    async def async_create(*args, **kwargs):  # type: ignore[no-untyped-def]
+        return mock_response
+
+    mock_openai.AsyncOpenAI.return_value.chat.completions.create = async_create
+
+    messages: List[LLMMessage] = [{"role": "user", "content": "Test"}]
+
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+    response = await llm.ainvoke(messages, response_format=_TestModelForOpenAI)
+
+    assert response.content == '{"value": "test"}'
+
+
+@pytest.mark.asyncio
+@patch("builtins.__import__")
+async def test_openai_llm_ainvoke_v2_with_json_schema_response_format(
+    mock_import: Mock,
+) -> None:
+    """Test V2 interface async invoke with JSON schema response_format."""
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    mock_response = MagicMock()
+    mock_response.choices = [
+        MagicMock(message=MagicMock(content='{"result": "success"}'))
+    ]
+
+    async def async_create(*args, **kwargs):  # type: ignore[no-untyped-def]
+        return mock_response
+
+    mock_openai.AsyncOpenAI.return_value.chat.completions.create = async_create
+
+    messages: List[LLMMessage] = [{"role": "user", "content": "Test"}]
+
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+    response = await llm.ainvoke(messages, response_format=_TEST_JSON_SCHEMA)
+
+    assert response.content == '{"result": "success"}'
+
+
+@patch("builtins.__import__")
+def test_openai_llm_close(mock_import: Mock) -> None:
+    mock_openai = get_mock_openai()
+    mock_import.side_effect = create_selective_import_mock(mock_openai)
+    mock_openai.AsyncOpenAI.return_value.close = AsyncMock()
+
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        llm.close()
+
+    mock_openai.OpenAI.return_value.close.assert_called_once()
+    mock_openai.AsyncOpenAI.return_value.close.assert_called_once()
+
+
+@pytest.mark.asyncio
+@patch("builtins.__import__")
+async def test_openai_llm_aclose(mock_import: Mock) -> None:
+    mock_openai = get_mock_openai()
+    mock_import.side_effect = create_selective_import_mock(mock_openai)
+    mock_openai.AsyncOpenAI.return_value.close = AsyncMock()
+
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        await llm.aclose()
+
+    mock_openai.OpenAI.return_value.close.assert_called_once()
+    mock_openai.AsyncOpenAI.return_value.close.assert_called_once()
+
+
+@pytest.mark.asyncio
+@patch("builtins.__import__")
+async def test_openai_llm_close_raises_in_async_context(mock_import: Mock) -> None:
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    llm = OpenAILLM(api_key="my key", model_name="gpt")
+
+    with pytest.raises(RuntimeError, match="async with"):
+        llm.close()
+
+
+# HTTP client tests
+
+
+@patch("builtins.__import__")
+def test_openai_llm_with_httpx_client(mock_import: Mock) -> None:
+    """Test that httpx.Client is forwarded only to the sync OpenAI client without warning."""
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    http_client = httpx.Client()
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        OpenAILLM(model_name="gpt", api_key="my key", http_client=http_client)
+
+    assert not any("Invalid http_client" in str(w.message) for w in caught)
+    _, sync_kwargs = mock_openai.OpenAI.call_args
+    assert sync_kwargs.get("http_client") is http_client
+    _, async_kwargs = mock_openai.AsyncOpenAI.call_args
+    assert async_kwargs.get("http_client") is None
+
+
+@patch("builtins.__import__")
+def test_openai_llm_with_httpx_async_client(mock_import: Mock) -> None:
+    """Test that httpx.AsyncClient is forwarded only to the async OpenAI client without warning."""
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    async_http_client = httpx.AsyncClient()
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        OpenAILLM(model_name="gpt", api_key="my key", http_client=async_http_client)
+
+    assert not any("Invalid http_client" in str(w.message) for w in caught)
+    _, sync_kwargs = mock_openai.OpenAI.call_args
+    assert sync_kwargs.get("http_client") is None
+    _, async_kwargs = mock_openai.AsyncOpenAI.call_args
+    assert async_kwargs.get("http_client") is async_http_client
+
+
+@patch("builtins.__import__")
+def test_openai_llm_no_http_client_no_warning(mock_import: Mock) -> None:
+    """Test that omitting http_client does not emit a warning."""
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        OpenAILLM(model_name="gpt", api_key="my key")
+
+    assert not any("Invalid http_client" in str(w.message) for w in caught)
+
+
+@patch("builtins.__import__")
+def test_openai_llm_with_invalid_http_client_warns(mock_import: Mock) -> None:
+    """Test that a non-None invalid http_client type emits a warning."""
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    with pytest.warns(UserWarning, match="Invalid http_client type"):
+        OpenAILLM(model_name="gpt", api_key="my key", http_client="not-a-client")
+
+
+@patch("builtins.__import__")
+def test_azure_openai_llm_with_httpx_client(mock_import: Mock) -> None:
+    """Test that httpx.Client is forwarded only to the sync AzureOpenAI client without warning."""
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    http_client = httpx.Client()
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        AzureOpenAILLM(
+            model_name="gpt",
+            azure_endpoint="https://test.openai.azure.com/",
+            api_key="my key",
+            api_version="version",
+            http_client=http_client,
+        )
+
+    assert not any("Invalid http_client" in str(w.message) for w in caught)
+    _, sync_kwargs = mock_openai.AzureOpenAI.call_args
+    assert sync_kwargs.get("http_client") is http_client
+    _, async_kwargs = mock_openai.AsyncAzureOpenAI.call_args
+    assert async_kwargs.get("http_client") is None
+
+
+@patch("builtins.__import__")
+def test_azure_openai_llm_with_httpx_async_client(mock_import: Mock) -> None:
+    """Test that httpx.AsyncClient is forwarded only to the async AzureOpenAI client without warning."""
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    async_http_client = httpx.AsyncClient()
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        AzureOpenAILLM(
+            model_name="gpt",
+            azure_endpoint="https://test.openai.azure.com/",
+            api_key="my key",
+            api_version="version",
+            http_client=async_http_client,
+        )
+
+    assert not any("Invalid http_client" in str(w.message) for w in caught)
+    _, sync_kwargs = mock_openai.AzureOpenAI.call_args
+    assert sync_kwargs.get("http_client") is None
+    _, async_kwargs = mock_openai.AsyncAzureOpenAI.call_args
+    assert async_kwargs.get("http_client") is async_http_client
+
+
+@patch("builtins.__import__")
+def test_openai_llm_with_default_aiohttp_client(mock_import: Mock) -> None:
+    """Test that DefaultAioHttpClient (subclass of httpx.AsyncClient) is forwarded to the async client.
+
+    DefaultAioHttpClient is a subclass of httpx.AsyncClient, so the isinstance check
+    already handles it without any special-casing.
+    """
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    class _FakeAioHttpClient(httpx.AsyncClient):
+        """Minimal stand-in for openai.DefaultAioHttpClient."""
+
+    aiohttp_client = _FakeAioHttpClient()
+
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        OpenAILLM(model_name="gpt", api_key="my key", http_client=aiohttp_client)
+
+    assert not any("Invalid http_client" in str(w.message) for w in caught)
+    _, sync_kwargs = mock_openai.OpenAI.call_args
+    assert sync_kwargs.get("http_client") is None
+    _, async_kwargs = mock_openai.AsyncOpenAI.call_args
+    assert async_kwargs.get("http_client") is aiohttp_client

참고/neo4j-graphrag-python-main/tests/unit/llm/test_rate_limit.py

@@ -0,0 +1,232 @@
+#  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.
+
+from typing import Any, Callable, Awaitable
+
+import pytest
+from unittest.mock import Mock
+from tenacity import RetryError
+
+from neo4j_graphrag.utils.rate_limit import (
+    RateLimitHandler,
+    NoOpRateLimitHandler,
+    DEFAULT_RATE_LIMIT_HANDLER,
+)
+from neo4j_graphrag.exceptions import RateLimitError
+
+
+def test_default_handler_retries_sync() -> None:
+    call_count = 0
+
+    def mock_func() -> None:
+        nonlocal call_count
+        call_count += 1
+        raise RateLimitError("Rate limit exceeded")
+
+    wrapped_func = DEFAULT_RATE_LIMIT_HANDLER.handle_sync(mock_func)
+
+    with pytest.raises(RetryError):
+        wrapped_func()
+
+    assert call_count == 3
+
+
+@pytest.mark.asyncio
+async def test_default_handler_retries_async() -> None:
+    call_count = 0
+
+    async def mock_func() -> None:
+        nonlocal call_count
+        call_count += 1
+        raise RateLimitError("Rate limit exceeded")
+
+    wrapped_func = DEFAULT_RATE_LIMIT_HANDLER.handle_async(mock_func)
+
+    with pytest.raises(RetryError):
+        await wrapped_func()
+
+    assert call_count == 3
+
+
+def test_other_errors_pass_through_sync() -> None:
+    call_count = 0
+
+    def mock_func() -> None:
+        nonlocal call_count
+        call_count += 1
+        raise ValueError("Some other error")
+
+    wrapped_func = DEFAULT_RATE_LIMIT_HANDLER.handle_sync(mock_func)
+
+    with pytest.raises(ValueError):
+        wrapped_func()
+
+    assert call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_other_errors_pass_through_async() -> None:
+    call_count = 0
+
+    async def mock_func() -> None:
+        nonlocal call_count
+        call_count += 1
+        raise ValueError("Some other error")
+
+    wrapped_func = DEFAULT_RATE_LIMIT_HANDLER.handle_async(mock_func)
+
+    with pytest.raises(ValueError):
+        await wrapped_func()
+
+    assert call_count == 1
+
+
+def test_noop_handler_sync() -> None:
+    def mock_func() -> str:
+        return "test result"
+
+    handler = NoOpRateLimitHandler()
+    wrapped_func = handler.handle_sync(mock_func)
+
+    assert wrapped_func() == "test result"
+    assert wrapped_func is mock_func
+
+
+@pytest.mark.asyncio
+async def test_noop_handler_async() -> None:
+    async def mock_func() -> str:
+        return "async test result"
+
+    handler = NoOpRateLimitHandler()
+    wrapped_func = handler.handle_async(mock_func)
+
+    assert await wrapped_func() == "async test result"
+    assert wrapped_func is mock_func
+
+
+def test_custom_handler_sync_retry_override() -> None:
+    call_count = 0
+
+    def mock_func() -> str:
+        nonlocal call_count
+        call_count += 1
+        if call_count == 1:
+            raise RateLimitError("Rate limit exceeded")
+        return "success after custom retry"
+
+    # Custom handler with single retry
+    def custom_handle_sync(func: Callable[[], Any]) -> Callable[[], Any]:
+        def wrapper() -> Any:
+            try:
+                return func()
+            except RateLimitError:
+                return func()  # Retry once
+
+        return wrapper
+
+    handler = Mock(spec=RateLimitHandler)
+    handler.handle_sync = custom_handle_sync
+
+    result = handler.handle_sync(mock_func)()
+    assert result == "success after custom retry"
+    assert call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_custom_handler_async_retry_override() -> None:
+    call_count = 0
+
+    async def mock_func() -> str:
+        nonlocal call_count
+        call_count += 1
+        if call_count == 1:
+            raise RateLimitError("Rate limit exceeded")
+        return "success after custom retry"
+
+    # Custom handler with single retry
+    def custom_handle_async(
+        func: Callable[[], Awaitable[Any]],
+    ) -> Callable[[], Awaitable[Any]]:
+        async def wrapper() -> Any:
+            try:
+                return await func()
+            except RateLimitError:
+                return await func()  # Retry once
+
+        return wrapper
+
+    handler = Mock(spec=RateLimitHandler)
+    handler.handle_async = custom_handle_async
+
+    result = await handler.handle_async(mock_func)()
+    assert result == "success after custom retry"
+    assert call_count == 2
+
+
+def test_deprecated_llm_rate_limit_module_import_warning() -> None:
+    import importlib
+    import sys
+
+    # Remove cached module to force re-import and trigger the module-level warning
+    sys.modules.pop("neo4j_graphrag.llm.rate_limit", None)
+
+    with pytest.warns(DeprecationWarning, match="neo4j_graphrag.utils.rate_limit"):
+        importlib.import_module("neo4j_graphrag.llm.rate_limit")
+
+
+def test_deprecated_llm_rate_limit_getattr_known_name() -> None:
+    import importlib
+    import sys
+
+    sys.modules.pop("neo4j_graphrag.llm.rate_limit", None)
+    with pytest.warns(DeprecationWarning):
+        mod = importlib.import_module("neo4j_graphrag.llm.rate_limit")
+
+    with pytest.warns(DeprecationWarning, match="neo4j_graphrag.utils.rate_limit"):
+        result = mod.RateLimitHandler
+
+    from neo4j_graphrag.utils.rate_limit import RateLimitHandler
+
+    assert result is RateLimitHandler
+
+
+def test_deprecated_llm_rate_limit_getattr_unknown_name() -> None:
+    import importlib
+    import sys
+
+    sys.modules.pop("neo4j_graphrag.llm.rate_limit", None)
+    with pytest.warns(DeprecationWarning):
+        mod = importlib.import_module("neo4j_graphrag.llm.rate_limit")
+
+    with pytest.raises(AttributeError):
+        _ = mod.NonExistent
+
+
+def test_deprecated_llm_init_getattr_known_name() -> None:
+    import neo4j_graphrag.llm as llm_module
+
+    with pytest.warns(DeprecationWarning, match="neo4j_graphrag.utils.rate_limit"):
+        result = llm_module.RateLimitHandler
+
+    from neo4j_graphrag.utils.rate_limit import RateLimitHandler
+
+    assert result is RateLimitHandler
+
+
+def test_deprecated_llm_init_getattr_unknown_name() -> None:
+    import neo4j_graphrag.llm as llm_module
+
+    with pytest.raises(AttributeError):
+        _ = llm_module.NonExistent

참고/neo4j-graphrag-python-main/tests/unit/llm/test_vertexai_llm.py

@@ -0,0 +1,649 @@
+#  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.
+from __future__ import annotations
+
+from typing import cast
+from typing import List
+
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
+import pytest
+from vertexai.generative_models import (
+    Content,
+    GenerationResponse,
+    Part,
+)
+
+from neo4j_graphrag.exceptions import LLMGenerationError
+from neo4j_graphrag.llm.types import ToolCallResponse
+from neo4j_graphrag.llm.vertexai_llm import VertexAILLM
+from neo4j_graphrag.tool import Tool
+from neo4j_graphrag.types import LLMMessage
+from neo4j_graphrag.utils.rate_limit import NoOpRateLimitHandler
+
+from pydantic import BaseModel, ConfigDict
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel", None)
+def test_vertexai_llm_missing_dependency() -> None:
+    with pytest.raises(ImportError):
+        VertexAILLM(model_name="gemini-1.5-flash-001")
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+def test_vertexai_llm_rate_limit_handler_is_set(
+    _GenerativeModelMock: MagicMock,
+) -> None:
+    custom_handler = MagicMock()
+    llm = VertexAILLM(
+        model_name="gemini-1.5-flash-001", rate_limit_handler=custom_handler
+    )
+    assert llm._rate_limit_handler is custom_handler
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+def test_vertexai_llm_default_rate_limit_handler_is_set(
+    _GenerativeModelMock: MagicMock,
+) -> None:
+    llm = VertexAILLM(model_name="gemini-1.5-flash-001")
+    assert hasattr(llm, "_rate_limit_handler")
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+def test_vertexai_invoke_happy_path(GenerativeModelMock: MagicMock) -> None:
+    model_name = "gemini-1.5-flash-001"
+    input_text = "may thy knife chip and shatter"
+    mock_response = Mock()
+    mock_response.text = "Return text"
+    mock_response.usage_metadata = None
+    mock_model = GenerativeModelMock.return_value
+    mock_model.generate_content.return_value = mock_response
+    model_params = {"temperature": 0.5}
+    llm = VertexAILLM(model_name, model_params)
+
+    response = llm.invoke(input_text)
+    assert response.content == "Return text"
+    GenerativeModelMock.assert_called_once_with(
+        model_name=model_name,
+        system_instruction=None,
+    )
+    last_call = mock_model.generate_content.call_args_list[0]
+    content = last_call.kwargs["contents"]
+    assert len(content) == 1
+    assert content[0].role == "user"
+    assert content[0].parts[0].text == input_text
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+@patch("neo4j_graphrag.llm.vertexai_llm.VertexAILLM.get_messages")
+def test_vertexai_invoke_with_system_instruction(
+    mock_get_messages: MagicMock,
+    GenerativeModelMock: MagicMock,
+) -> None:
+    system_instruction = "You are a helpful assistant."
+    model_name = "gemini-1.5-flash-001"
+    input_text = "may thy knife chip and shatter"
+    mock_response = Mock()
+    mock_response.text = "Return text"
+    mock_response.usage_metadata = None
+    mock_model = GenerativeModelMock.return_value
+    mock_model.generate_content.return_value = mock_response
+
+    mock_get_messages.return_value = [{"text": "some text"}]
+
+    model_params = {"temperature": 0.5}
+    llm = VertexAILLM(model_name, model_params)
+
+    response = llm.invoke(input_text, system_instruction=system_instruction)
+    assert response.content == "Return text"
+    GenerativeModelMock.assert_called_once_with(
+        model_name=model_name,
+        system_instruction=system_instruction,
+    )
+    mock_model.generate_content.assert_called_once_with(
+        contents=[{"text": "some text"}]
+    )
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+def test_vertexai_invoke_with_message_history_and_system_instruction(
+    GenerativeModelMock: MagicMock,
+) -> None:
+    system_instruction = "You are a helpful assistant."
+    model_name = "gemini-1.5-flash-001"
+    mock_response = Mock()
+    mock_response.text = "Return text"
+    mock_response.usage_metadata = None
+    mock_model = GenerativeModelMock.return_value
+    mock_model.generate_content.return_value = mock_response
+    model_params = {"temperature": 0.5}
+    llm = VertexAILLM(model_name, model_params)
+
+    message_history = [
+        {"role": "user", "content": "When does the sun come up in the summer?"},
+        {"role": "assistant", "content": "Usually around 6am."},
+    ]
+    question = "What about next season?"
+
+    response = llm.invoke(
+        question,
+        message_history,  # type: ignore
+        system_instruction=system_instruction,
+    )
+    assert response.content == "Return text"
+    GenerativeModelMock.assert_called_once_with(
+        model_name=model_name,
+        system_instruction=system_instruction,
+    )
+    last_call = mock_model.generate_content.call_args_list[0]
+    content = last_call.kwargs["contents"]
+    assert len(content) == 3  # question + 2 messages in history
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+def test_vertexai_get_messages(GenerativeModelMock: MagicMock) -> None:
+    model_name = "gemini-1.5-flash-001"
+    question = "When does it set?"
+    message_history: list[LLMMessage] = [
+        {"role": "user", "content": "When does the sun come up in the summer?"},
+        {"role": "assistant", "content": "Usually around 6am."},
+        {"role": "user", "content": "What about next season?"},
+        {"role": "assistant", "content": "Around 8am."},
+    ]
+    expected_response = [
+        Content(
+            role="user",
+            parts=[Part.from_text("When does the sun come up in the summer?")],
+        ),
+        Content(role="model", parts=[Part.from_text("Usually around 6am.")]),
+        Content(role="user", parts=[Part.from_text("What about next season?")]),
+        Content(role="model", parts=[Part.from_text("Around 8am.")]),
+        Content(role="user", parts=[Part.from_text("When does it set?")]),
+    ]
+
+    llm = VertexAILLM(model_name=model_name)
+    response = llm.get_messages(question, message_history)
+
+    GenerativeModelMock.assert_not_called()
+    assert len(response) == len(expected_response)
+    for actual, expected in zip(response, expected_response):
+        assert actual.role == expected.role
+        assert actual.parts[0].text == expected.parts[0].text
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+def test_vertexai_get_messages_validation_error(
+    _GenerativeModelMock: MagicMock,
+) -> None:
+    system_instruction = "You are a helpful assistant."
+    model_name = "gemini-1.5-flash-001"
+    question = "hi!"
+    message_history = [
+        {"role": "model", "content": "hello!"},
+    ]
+
+    llm = VertexAILLM(model_name=model_name, system_instruction=system_instruction)
+    with pytest.raises(LLMGenerationError) as exc_info:
+        llm.invoke(question, cast(list[LLMMessage], message_history))
+    assert "Input should be 'user', 'assistant' or 'system" in str(exc_info.value)
+
+
+@pytest.mark.asyncio
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+@patch("neo4j_graphrag.llm.vertexai_llm.VertexAILLM.get_messages")
+async def test_vertexai_ainvoke_happy_path(
+    mock_get_messages: Mock, GenerativeModelMock: MagicMock
+) -> None:
+    mock_response = AsyncMock()
+    mock_response.text = "Return text"
+    mock_model = GenerativeModelMock.return_value
+    mock_model.generate_content_async = AsyncMock(return_value=mock_response)
+    mock_get_messages.return_value = [{"text": "Return text"}]
+    model_params = {"temperature": 0.5}
+    llm = VertexAILLM("gemini-1.5-flash-001", model_params)
+    input_text = "may thy knife chip and shatter"
+    response = await llm.ainvoke(input_text)
+    print(f"Response: {response}")
+    assert response.content == "Return text"
+    mock_model.generate_content_async.assert_awaited_once_with(
+        contents=[{"text": "Return text"}]
+    )
+
+
+def test_vertexai_get_llm_tools(test_tool: Tool) -> None:
+    llm = VertexAILLM(model_name="gemini")
+    tools = llm._get_llm_tools(tools=[test_tool])
+    assert tools is not None
+    assert len(tools) == 1
+    tool = tools[0]
+    tool_dict = tool.to_dict()
+    assert len(tool_dict["function_declarations"]) == 1
+    assert tool_dict["function_declarations"][0]["name"] == "test_tool"
+    assert tool_dict["function_declarations"][0]["description"] == "A test tool"
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.VertexAILLM._parse_tool_response")
+@patch("neo4j_graphrag.llm.vertexai_llm.VertexAILLM._call_llm")
+def test_vertexai_invoke_with_tools(
+    mock_call_llm: Mock,
+    mock_parse_tool: Mock,
+    test_tool: Tool,
+) -> None:
+    # Mock the model call response
+    tool_call_mock = MagicMock()
+    tool_call_mock.name = "function"
+    tool_call_mock.args = {}
+    mock_call_llm.return_value = MagicMock(
+        candidates=[MagicMock(function_calls=[tool_call_mock])]
+    )
+    mock_parse_tool.return_value = ToolCallResponse(tool_calls=[])
+
+    llm = VertexAILLM(model_name="gemini")
+    tools = [test_tool]
+
+    res = llm.invoke_with_tools("my text", tools)
+    mock_call_llm.assert_called_once_with(
+        "my text",
+        message_history=None,
+        system_instruction=None,
+        tools=tools,
+    )
+    mock_parse_tool.assert_called_once()
+    assert isinstance(res, ToolCallResponse)
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.VertexAILLM._get_model")
+def test_vertexai_call_llm_with_tools(mock_model: Mock, test_tool: Tool) -> None:
+    # Mock the generation response
+    mock_generate_content = mock_model.return_value.generate_content
+    mock_generate_content.return_value = MagicMock(
+        spec=GenerationResponse,
+    )
+
+    llm = VertexAILLM(model_name="gemini")
+    tools = [test_tool]
+
+    with patch.object(llm, "_get_llm_tools", return_value=["my tools"]):
+        res = llm._call_llm("my text", tools=tools)
+        assert isinstance(res, GenerationResponse)
+
+        mock_model.assert_called_once_with(
+            system_instruction=None,
+        )
+        calls = mock_generate_content.call_args_list
+        assert len(calls) == 1
+        assert calls[0][1]["tools"] == ["my tools"]
+        assert calls[0][1]["tool_config"] is not None
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.VertexAILLM._parse_tool_response")
+@patch("neo4j_graphrag.llm.vertexai_llm.VertexAILLM._call_llm")
+def test_vertexai_ainvoke_with_tools(
+    mock_call_llm: Mock,
+    mock_parse_tool: Mock,
+    test_tool: Tool,
+) -> None:
+    # Mock the model call response
+    tool_call_mock = MagicMock()
+    tool_call_mock.name = "function"
+    tool_call_mock.args = {}
+    mock_call_llm.return_value = AsyncMock(
+        return_value=MagicMock(candidates=[MagicMock(function_calls=[tool_call_mock])])
+    )
+    mock_parse_tool.return_value = ToolCallResponse(tool_calls=[])
+
+    llm = VertexAILLM(model_name="gemini")
+    tools = [test_tool]
+
+    res = llm.invoke_with_tools("my text", tools)
+    mock_call_llm.assert_called_once_with(
+        "my text",
+        message_history=None,
+        system_instruction=None,
+        tools=tools,
+    )
+    mock_parse_tool.assert_called_once()
+    assert isinstance(res, ToolCallResponse)
+
+
+@pytest.mark.asyncio
+@patch("neo4j_graphrag.llm.vertexai_llm.VertexAILLM._get_model")
+async def test_vertexai_acall_llm_with_tools(mock_model: Mock, test_tool: Tool) -> None:
+    # Mock the generation response
+    mock_model.return_value = AsyncMock(
+        generate_content_async=AsyncMock(
+            return_value=MagicMock(
+                spec=GenerationResponse,
+            )
+        )
+    )
+
+    llm = VertexAILLM(model_name="gemini")
+    tools = [test_tool]
+
+    res = await llm._acall_llm("my text", tools=tools)
+    mock_model.assert_called_once_with(
+        system_instruction=None,
+    )
+    assert isinstance(res, GenerationResponse)
+
+
+# LLM Interface V2 Tests
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+def test_vertexai_invoke_v2_happy_path(GenerativeModelMock: MagicMock) -> None:
+    """Test V2 interface invoke method with List[LLMMessage] input."""
+    model_name = "gemini-1.5-flash-001"
+    messages: List[LLMMessage] = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "What is the capital of France?"},
+    ]
+    mock_response = Mock()
+    mock_response.text = "Paris is the capital of France."
+    mock_response.usage_metadata = None
+    mock_model = GenerativeModelMock.return_value
+    mock_model.generate_content.return_value = mock_response
+
+    llm = VertexAILLM(model_name=model_name)
+    response = llm.invoke(messages)
+
+    assert response.content == "Paris is the capital of France."
+    GenerativeModelMock.assert_called_once_with(
+        model_name=model_name,
+        system_instruction="You are a helpful assistant.",
+    )
+    mock_model.generate_content.assert_called_once()
+    call_args = mock_model.generate_content.call_args
+    contents = call_args.kwargs["contents"]
+    assert len(contents) == 1  # Only user message after system is extracted
+    assert contents[0].role == "user"
+    assert contents[0].parts[0].text == "What is the capital of France?"
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+def test_vertexai_invoke_v2_with_conversation_history(
+    GenerativeModelMock: MagicMock,
+) -> None:
+    """Test V2 interface invoke with conversation history."""
+    model_name = "gemini-1.5-flash-001"
+    messages: List[LLMMessage] = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "What is the capital of France?"},
+        {"role": "assistant", "content": "Paris is the capital of France."},
+        {"role": "user", "content": "What about Germany?"},
+    ]
+    mock_response = Mock()
+    mock_response.text = "Berlin is the capital of Germany."
+    mock_response.usage_metadata = None
+    mock_model = GenerativeModelMock.return_value
+    mock_model.generate_content.return_value = mock_response
+
+    llm = VertexAILLM(model_name=model_name)
+    response = llm.invoke(messages)
+
+    assert response.content == "Berlin is the capital of Germany."
+    GenerativeModelMock.assert_called_once_with(
+        model_name=model_name,
+        system_instruction="You are a helpful assistant.",
+    )
+    call_args = mock_model.generate_content.call_args
+    contents = call_args.kwargs["contents"]
+    assert len(contents) == 3  # user -> assistant -> user
+    assert contents[0].role == "user"
+    assert contents[0].parts[0].text == "What is the capital of France?"
+    assert contents[1].role == "model"
+    assert contents[1].parts[0].text == "Paris is the capital of France."
+    assert contents[2].role == "user"
+    assert contents[2].parts[0].text == "What about Germany?"
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+def test_vertexai_invoke_v2_no_system_message(GenerativeModelMock: MagicMock) -> None:
+    """Test V2 interface invoke without system message."""
+    model_name = "gemini-1.5-flash-001"
+    messages: List[LLMMessage] = [
+        {"role": "user", "content": "Hello, how are you?"},
+    ]
+    mock_response = Mock()
+    mock_response.text = "I'm doing well, thank you!"
+    mock_response.usage_metadata = None
+    mock_model = GenerativeModelMock.return_value
+    mock_model.generate_content.return_value = mock_response
+
+    llm = VertexAILLM(model_name=model_name)
+    response = llm.invoke(messages)
+
+    assert response.content == "I'm doing well, thank you!"
+    GenerativeModelMock.assert_called_once_with(
+        model_name=model_name,
+        system_instruction=None,  # No system instruction should be used
+    )
+
+
+@pytest.mark.asyncio
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+async def test_vertexai_ainvoke_v2_happy_path(GenerativeModelMock: MagicMock) -> None:
+    """Test V2 interface async invoke method with List[LLMMessage] input."""
+    model_name = "gemini-1.5-flash-001"
+    messages: List[LLMMessage] = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "What is 2+2?"},
+    ]
+    mock_response = AsyncMock()
+    mock_response.text = "2+2 equals 4."
+    mock_model = GenerativeModelMock.return_value
+    mock_model.generate_content_async = AsyncMock(return_value=mock_response)
+
+    llm = VertexAILLM(model_name=model_name)
+    response = await llm.ainvoke(messages)
+
+    assert response.content == "2+2 equals 4."
+    GenerativeModelMock.assert_called_once_with(
+        model_name=model_name,
+        system_instruction="You are a helpful assistant.",
+    )
+    mock_model.generate_content_async.assert_awaited_once()
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+def test_vertexai_invoke_v2_validation_error(_GenerativeModelMock: MagicMock) -> None:
+    """Test V2 interface invoke with invalid role raises error."""
+    model_name = "gemini-1.5-flash-001"
+    messages: List[LLMMessage] = [
+        {"role": "invalid_role", "content": "This should fail."},  # type: ignore[typeddict-item]
+    ]
+
+    llm = VertexAILLM(model_name=model_name)
+
+    with pytest.raises(ValueError) as exc_info:
+        llm.invoke(messages)
+    assert "Unknown role: invalid_role" in str(exc_info.value)
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+def test_vertexai_get_brand_new_messages_system_instruction_override(
+    _GenerativeModelMock: MagicMock,
+) -> None:
+    """Test that system instruction in messages overrides class-level system instruction."""
+    model_name = "gemini-1.5-flash-001"
+    class_system_instruction = "You are a class-level assistant."
+    messages: List[LLMMessage] = [
+        {"role": "system", "content": "You are a message-level assistant."},
+        {"role": "user", "content": "Hello"},
+    ]
+
+    llm = VertexAILLM(
+        model_name=model_name, system_instruction=class_system_instruction
+    )
+    system_instruction, contents = llm.get_messages_v2(messages)
+
+    assert system_instruction == "You are a message-level assistant."
+    assert len(contents) == 1  # Only user message should remain
+    assert contents[0].role == "user"
+    assert contents[0].parts[0].text == "Hello"
+
+
+class _TestModelForVertexAI(BaseModel):
+    """Test model for structured output tests."""
+
+    model_config = ConfigDict(extra="forbid")
+    name: str
+    age: int
+
+
+_TEST_JSON_SCHEMA = {
+    "type": "object",
+    "properties": {"result": {"type": "string"}},
+    "required": ["result"],
+}
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+def test_vertexai_invoke_v2_with_pydantic_response_format(
+    GenerativeModelMock: MagicMock,
+) -> None:
+    """Test V2 interface with Pydantic model as response_format."""
+
+    model_name = "gemini-2.5-flash"
+    messages: List[LLMMessage] = [
+        {"role": "user", "content": "Extract person info"},
+    ]
+    mock_response = Mock()
+    mock_response.text = '{"name": "John", "age": 30}'
+    mock_response.usage_metadata = None
+    mock_model = GenerativeModelMock.return_value
+    mock_model.generate_content.return_value = mock_response
+
+    llm = VertexAILLM(model_name=model_name)
+    response = llm.invoke(messages, response_format=_TestModelForVertexAI)
+
+    assert response.content == '{"name": "John", "age": 30}'
+
+    # Verify the method was called with generation_config
+    mock_model.generate_content.assert_called_once()
+    call_args = mock_model.generate_content.call_args.kwargs
+    assert "generation_config" in call_args
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+def test_vertexai_invoke_v2_with_json_schema_response_format(
+    GenerativeModelMock: MagicMock,
+) -> None:
+    """Test V2 interface with JSON schema dict as response_format."""
+    model_name = "gemini-2.5-flash"
+    messages: List[LLMMessage] = [
+        {"role": "user", "content": "Test"},
+    ]
+    mock_response = Mock()
+    mock_response.text = '{"result": "success"}'
+    mock_response.usage_metadata = None
+    mock_model = GenerativeModelMock.return_value
+    mock_model.generate_content.return_value = mock_response
+
+    llm = VertexAILLM(model_name=model_name)
+    response = llm.invoke(messages, response_format=_TEST_JSON_SCHEMA)
+
+    assert response.content == '{"result": "success"}'
+
+    # Verify the method was called with generation_config
+    mock_model.generate_content.assert_called_once()
+    call_args = mock_model.generate_content.call_args.kwargs
+    assert "generation_config" in call_args
+
+
+@pytest.mark.asyncio
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+async def test_vertexai_ainvoke_v2_with_pydantic_response_format(
+    GenerativeModelMock: MagicMock,
+) -> None:
+    """Test V2 interface async invoke with Pydantic response_format."""
+
+    model_name = "gemini-2.5-flash"
+    messages: List[LLMMessage] = [{"role": "user", "content": "Test"}]
+    mock_response = AsyncMock()
+    mock_response.text = '{"value": "test"}'
+    mock_model = GenerativeModelMock.return_value
+    mock_model.generate_content_async = AsyncMock(return_value=mock_response)
+
+    llm = VertexAILLM(model_name=model_name)
+    response = await llm.ainvoke(messages, response_format=_TestModelForVertexAI)
+
+    assert response.content == '{"value": "test"}'
+
+    # Verify generation_config has response_schema
+    call_args = mock_model.generate_content_async.call_args.kwargs
+    assert "generation_config" in call_args
+
+
+@pytest.mark.asyncio
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+async def test_vertexai_ainvoke_v2_with_json_schema_response_format(
+    GenerativeModelMock: MagicMock,
+) -> None:
+    """Test V2 interface async invoke with JSON schema response_format."""
+    model_name = "gemini-2.5-flash"
+    messages: List[LLMMessage] = [{"role": "user", "content": "Test"}]
+    mock_response = AsyncMock()
+    mock_response.text = '{"result": "success"}'
+    mock_model = GenerativeModelMock.return_value
+    mock_model.generate_content_async = AsyncMock(return_value=mock_response)
+
+    llm = VertexAILLM(model_name=model_name)
+    response = await llm.ainvoke(messages, response_format=_TEST_JSON_SCHEMA)
+
+    assert response.content == '{"result": "success"}'
+
+    # Verify generation_config has response_schema
+    call_args = mock_model.generate_content_async.call_args.kwargs
+    assert "generation_config" in call_args
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
+def test_vertexai_invoke_v2_rate_limit_handler_called(
+    GenerativeModelMock: MagicMock,
+) -> None:
+    """Test that the rate limit handler is invoked on the V2 (List[LLMMessage]) path."""
+    messages: List[LLMMessage] = [{"role": "user", "content": "Hello"}]
+    mock_response = Mock()
+    mock_response.text = "Hi there!"
+    mock_response.usage_metadata = None
+    mock_model = GenerativeModelMock.return_value
+    mock_model.generate_content.return_value = mock_response
+
+    spy_handler = MagicMock(wraps=NoOpRateLimitHandler())
+    llm = VertexAILLM(model_name="gemini-1.5-flash-001", 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.vertexai_llm.GenerativeModel")
+async def test_vertexai_ainvoke_v2_rate_limit_handler_called(
+    GenerativeModelMock: MagicMock,
+) -> None:
+    """Test that the rate limit handler is invoked on the async V2 (List[LLMMessage]) path."""
+    messages: List[LLMMessage] = [{"role": "user", "content": "Hello"}]
+    mock_response = AsyncMock()
+    mock_response.text = "Hi there!"
+    mock_model = GenerativeModelMock.return_value
+    mock_model.generate_content_async = AsyncMock(return_value=mock_response)
+
+    spy_handler = MagicMock(wraps=NoOpRateLimitHandler())
+    llm = VertexAILLM(model_name="gemini-1.5-flash-001", rate_limit_handler=spy_handler)
+    response = await llm.ainvoke(messages)
+
+    assert response.content == "Hi there!"
+    spy_handler.handle_async.assert_called_once()