Update Ollama/Anthropic

Author: stellasia

tests/unit/llm/test_anthropic_llm.py

@@ -19,9 +19,11 @@
 
 import anthropic
 import pytest
-from neo4j_graphrag.exceptions import LLMGenerationError
+from anthropic import NOT_GIVEN, NotGiven
+
+from neo4j_graphrag.llm import LLMResponse
 from neo4j_graphrag.llm.anthropic_llm import AnthropicLLM
-from neo4j_graphrag.llm.types import LLMResponse
+from neo4j_graphrag.types import LLMMessage
 
 
 @pytest.fixture
@@ -40,147 +42,82 @@ def test_anthropic_llm_missing_dependency(mock_import: Mock) -> None:
         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"
-    llm.client.messages.create.assert_called_once_with(  # type: ignore
-        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,
-    )
+def test_anthropic_llm_get_messages_with_system_instructions() -> None:
+    llm = AnthropicLLM(api_key="my key", model_name="claude")
     message_history = [
-        {"role": "user", "content": "When does the sun come up in the summer?"},
-        {"role": "assistant", "content": "Usually around 6am."},
+        LLMMessage(**{"role": "system", "content": "do something"}),
+        LLMMessage(
+            **{"role": "user", "content": "When does the sun come up in the summer?"}
+        ),
+        LLMMessage(**{"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})
-    llm.client.messages.create.assert_called_once_with(  # type: ignore[attr-defined]
-        messages=message_history,
-        model="claude-3-opus-20240229",
-        system=anthropic.NOT_GIVEN,
-        **model_params,
-    )
 
+    system_instruction, messages = llm.get_messages(message_history)
+    assert isinstance(system_instruction, str)
+    assert system_instruction == "do something"
+    assert isinstance(messages, list)
+    assert len(messages) == 2  # exclude system instruction
+    for actual, expected in zip(messages, message_history[1:]):
+        assert isinstance(actual, dict)
+        assert actual["role"] == expected["role"]
+        assert actual["content"] == expected["content"]
 
-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 = [{"role": "user", "content": question}]
-    llm.client.messages.create.assert_called_with(  # type: ignore[attr-defined]
-        model="claude-3-opus-20240229",
-        system=system_instruction,
-        messages=messages,
-        **model_params,
-    )
+def test_anthropic_llm_get_messages_without_system_instructions() -> None:
+    llm = AnthropicLLM(api_key="my key", model_name="claude")
+    message_history = [
+        LLMMessage(
+            **{"role": "user", "content": "When does the sun come up in the summer?"}
+        ),
+        LLMMessage(**{"role": "assistant", "content": "Usually around 6am."}),
+    ]
 
-    assert llm.client.messages.create.call_count == 1  # type: ignore
+    system_instruction, messages = llm.get_messages(message_history)
+    assert isinstance(system_instruction, NotGiven)
+    assert system_instruction == NOT_GIVEN
+    assert isinstance(messages, list)
+    assert len(messages) == 2
+    for actual, expected in zip(messages, message_history):
+        assert isinstance(actual, dict)
+        assert actual["role"] == expected["role"]
+        assert actual["content"] == expected["content"]
 
 
-def test_anthropic_invoke_with_message_history_and_system_instruction(
-    mock_anthropic: Mock,
-) -> None:
+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")]
     )
+    mock_anthropic.types.MessageParam.return_value = {"role": "user", "content": "hi"}
     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
+    llm = AnthropicLLM("claude-3-opus-20240229", model_params=model_params)
+    input_text = "may thy knife chip and shatter"
+    response = llm.invoke(input_text)
     assert isinstance(response, LLMResponse)
     assert response.content == "generated text"
-    message_history.append({"role": "user", "content": question})
-    llm.client.messages.create.assert_called_with(  # type: ignore[attr-defined]
+    llm.client.messages.create.assert_called_once_with(  # type: ignore
+        messages=[{"role": "user", "content": "hi"}],
         model="claude-3-opus-20240229",
-        system=system_instruction,
-        messages=message_history,
+        system=anthropic.NOT_GIVEN,
         **model_params,
     )
 
-    assert llm.client.messages.create.call_count == 1  # type: ignore
-
-
-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)
+    mock_anthropic.types.MessageParam.return_value = {"role": "user", "content": "hi"}
     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 isinstance(response, LLMResponse)
     assert response.content == "Return text"
     llm.async_client.messages.create.assert_awaited_once_with(  # type: ignore
         model="claude-3-opus-20240229",
         system=anthropic.NOT_GIVEN,
-        messages=[{"role": "user", "content": input_text}],
+        messages=[{"role": "user", "content": "hi"}],
         **model_params,
     )

tests/unit/llm/test_ollama_llm.py

@@ -62,10 +62,8 @@ def test_ollama_llm_happy_path_deprecated_options(mock_import: Mock) -> None:
     assert res.content == "ollama chat response"
     llm.client.chat.assert_called_once_with(  # type: ignore[attr-defined]
         model=model,
-        messages=[
-            {"role": "user", "content": "test"}
-        ],
-        options={"temperature": 0.3}
+        messages=[{"role": "user", "content": "test"}],
+        options={"temperature": 0.3},
     )