fix mypy errors

Author: oskarhane

src/neo4j_graphrag/embeddings/cohere.py

@@ -21,7 +21,7 @@
 try:
     import cohere
 except ImportError:
-    cohere = None
+    cohere = None  # type: ignore[assignment]
 
 
 class CohereEmbeddings(Embedder):

src/neo4j_graphrag/embeddings/mistral.py

@@ -24,7 +24,10 @@
 try:
     from mistralai import Mistral
 except ImportError:
-    Mistral = None
+    # Define placeholder type for type checking
+    class Mistral:  # type: ignore
+        pass
+    Mistral = None  # type: ignore
 
 
 class MistralAIEmbeddings(Embedder):

src/neo4j_graphrag/llm/anthropic_llm.py

@@ -91,7 +91,7 @@ def get_messages(
                 raise LLMGenerationError(e.errors()) from e
             messages.extend(cast(Iterable[dict[str, Any]], message_history))
         messages.append(UserMessage(content=input).model_dump())
-        return messages
+        return messages  # type: ignore[return-value]
 
     def invoke(
         self,

src/neo4j_graphrag/llm/cohere_llm.py

@@ -94,7 +94,7 @@ def get_messages(
                 raise LLMGenerationError(e.errors()) from e
             messages.extend(cast(Iterable[dict[str, Any]], message_history))
         messages.append(UserMessage(content=input).model_dump())
-        return messages
+        return messages  # type: ignore[return-value]
 
     def invoke(
         self,

src/neo4j_graphrag/llm/mistralai_llm.py

@@ -35,8 +35,13 @@
     from mistralai import Messages, Mistral
     from mistralai.models.sdkerror import SDKError
 except ImportError:
-    Mistral = None
-    SDKError = None
+    # Define placeholder types for type checking
+    class Mistral:  # type: ignore
+        pass
+    class SDKError(Exception):  # type: ignore
+        pass
+    Mistral = None  # type: ignore
+    SDKError = None  # type: ignore
 
 
 class MistralAILLM(LLMInterface):

src/neo4j_graphrag/llm/ollama_llm.py

@@ -76,7 +76,7 @@ def get_messages(
                 raise LLMGenerationError(e.errors()) from e
             messages.extend(cast(Iterable[dict[str, Any]], message_history))
         messages.append(UserMessage(content=input).model_dump())
-        return messages
+        return messages  # type: ignore[return-value]
 
     def invoke(
         self,

tests/unit/llm/test_anthropic_llm.py

@@ -49,7 +49,7 @@ def test_anthropic_invoke_happy_path(mock_anthropic: Mock) -> None:
     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(
+    llm.client.messages.create.assert_called_once_with(  # type: ignore[attr-defined]
         messages=[{"role": "user", "content": input_text}],
         model="claude-3-opus-20240229",
         system=anthropic.NOT_GIVEN,
@@ -81,7 +81,7 @@ def test_anthropic_invoke_with_message_history_happy_path(mock_anthropic: Mock)
     response = llm.invoke(question, message_history)
     assert response.content == "generated text"
     message_history.add_message(LLMMessage(role="user", content=question))
-    llm.client.messages.create.assert_called_once_with(
+    llm.client.messages.create.assert_called_once_with(  # type: ignore[attr-defined]
         messages=message_history,
         model="claude-3-opus-20240229",
         system=anthropic.NOT_GIVEN,
@@ -107,14 +107,14 @@ def test_anthropic_invoke_with_system_instruction(
     assert isinstance(response, LLMResponse)
     assert response.content == "generated text"
     messages = [{"role": "user", "content": question}]
-    llm.client.messages.create.assert_called_with(
+    llm.client.messages.create.assert_called_with(  # type: ignore[attr-defined]
         model="claude-3-opus-20240229",
         system=system_instruction,
         messages=messages,
         **model_params,
     )
 
-    assert llm.client.messages.create.call_count == 1
+    assert llm.client.messages.create.call_count == 1  # type: ignore[attr-defined]
 
 
 def test_anthropic_invoke_with_message_history_and_system_instruction(
@@ -145,14 +145,14 @@ def test_anthropic_invoke_with_message_history_and_system_instruction(
     assert isinstance(response, LLMResponse)
     assert response.content == "generated text"
     message_history.add_message(LLMMessage(role="user", content=question))
-    llm.client.messages.create.assert_called_with(
+    llm.client.messages.create.assert_called_with(  # type: ignore[attr-defined]
         model="claude-3-opus-20240229",
         system=system_instruction,
         messages=message_history,
         **model_params,
     )
 
-    assert llm.client.messages.create.call_count == 1
+    assert llm.client.messages.create.call_count == 1  # type: ignore[attr-defined]
 
 
 def test_anthropic_invoke_with_message_history_validation_error(
@@ -190,7 +190,7 @@ async def test_anthropic_ainvoke_happy_path(mock_anthropic: Mock) -> None:
     input_text = "may thy knife chip and shatter"
     response = await llm.ainvoke(input_text)
     assert response.content == "Return text"
-    llm.async_client.messages.create.assert_awaited_once_with(
+    llm.async_client.messages.create.assert_awaited_once_with(  # type: ignore[attr-defined]
         model="claude-3-opus-20240229",
         system=anthropic.NOT_GIVEN,
         messages=[{"role": "user", "content": input_text}],

tests/unit/llm/test_mistralai_llm.py

@@ -71,7 +71,7 @@ def test_mistralai_llm_invoke_with_message_history(mock_mistral: Mock) -> None:
     messages = [{"role": "system", "content": system_instruction}]
     messages.extend(message_history)
     messages.append({"role": "user", "content": question})
-    llm.client.chat.complete.assert_called_once_with(
+    llm.client.chat.complete.assert_called_once_with(  # type: ignore[attr-defined]
         messages=messages,
         model=model,
     )
@@ -103,12 +103,12 @@ def test_mistralai_llm_invoke_with_message_history_and_system_instruction(
     messages = [{"role": "system", "content": system_instruction}]
     messages.extend(message_history)
     messages.append({"role": "user", "content": question})
-    llm.client.chat.complete.assert_called_once_with(
+    llm.client.chat.complete.assert_called_once_with(  # type: ignore[attr-defined]
         messages=messages,
         model=model,
     )
 
-    assert llm.client.chat.complete.call_count == 1
+    assert llm.client.chat.complete.call_count == 1  # type: ignore[attr-defined]
 
 
 @patch("neo4j_graphrag.llm.mistralai_llm.Mistral")

tests/unit/llm/test_ollama_llm.py

@@ -55,7 +55,7 @@ def test_ollama_llm_happy_path(mock_import: Mock) -> None:
     messages = [
         {"role": "user", "content": question},
     ]
-    llm.client.chat.assert_called_once_with(
+    llm.client.chat.assert_called_once_with(  # type: ignore[attr-defined]
         model=model, messages=messages, options=model_params
     )
 
@@ -80,7 +80,7 @@ def test_ollama_invoke_with_system_instruction_happy_path(mock_import: Mock) ->
     assert response.content == "ollama chat response"
     messages = [{"role": "system", "content": system_instruction}]
     messages.append({"role": "user", "content": question})
-    llm.client.chat.assert_called_once_with(
+    llm.client.chat.assert_called_once_with(  # type: ignore[attr-defined]
         model=model, messages=messages, options=model_params
     )
 
@@ -108,7 +108,7 @@ def test_ollama_invoke_with_message_history_happy_path(mock_import: Mock) -> Non
     assert response.content == "ollama chat response"
     messages = [m for m in message_history]
     messages.append({"role": "user", "content": question})
-    llm.client.chat.assert_called_once_with(
+    llm.client.chat.assert_called_once_with(  # type: ignore[attr-defined]
         model=model, messages=messages, options=model_params
     )
 
@@ -144,10 +144,10 @@ def test_ollama_invoke_with_message_history_and_system_instruction(
     messages = [{"role": "system", "content": system_instruction}]
     messages.extend(message_history)
     messages.append({"role": "user", "content": question})
-    llm.client.chat.assert_called_once_with(
+    llm.client.chat.assert_called_once_with(  # type: ignore[attr-defined]
         model=model, messages=messages, options=model_params
     )
-    assert llm.client.chat.call_count == 1
+    assert llm.client.chat.call_count == 1  # type: ignore[attr-defined]
 
 
 @patch("builtins.__import__")