Try to remove type-ignore comments

Author: stellasia

src/neo4j_graphrag/embeddings/cohere.py

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

src/neo4j_graphrag/embeddings/mistral.py

@@ -24,7 +24,7 @@
 try:
     from mistralai import Mistral
 except ImportError:
-    Mistral = None  # type: ignore
+    Mistral = None
 
 
 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  # type: ignore
+        return messages
 
     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  # type: ignore
+        return messages
 
     def invoke(
         self,

src/neo4j_graphrag/llm/mistralai_llm.py

@@ -35,8 +35,8 @@
     from mistralai import Messages, Mistral
     from mistralai.models.sdkerror import SDKError
 except ImportError:
-    Mistral = None  # type: ignore
-    SDKError = None  # type: ignore
+    Mistral = None
+    SDKError = None
 
 
 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  # type: ignore
+        return messages
 
     def invoke(
         self,

src/neo4j_graphrag/llm/openai_llm.py

@@ -85,7 +85,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  # type: ignore
+        return messages
 
     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(  # type: ignore
+    llm.client.messages.create.assert_called_once_with(
         messages=[{"role": "user", "content": input_text}],
         model="claude-3-opus-20240229",
         system=anthropic.NOT_GIVEN,
@@ -75,7 +75,7 @@ def test_anthropic_invoke_with_message_history_happy_path(mock_anthropic: Mock)
     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]
+    llm.client.messages.create.assert_called_once_with(
         messages=message_history,
         model="claude-3-opus-20240229",
         system=anthropic.NOT_GIVEN,
@@ -101,14 +101,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(  # type: ignore[attr-defined]
+    llm.client.messages.create.assert_called_with(
         model="claude-3-opus-20240229",
         system=system_instruction,
         messages=messages,
         **model_params,
     )
 
-    assert llm.client.messages.create.call_count == 1  # type: ignore
+    assert llm.client.messages.create.call_count == 1
 
 
 def test_anthropic_invoke_with_message_history_and_system_instruction(
@@ -133,14 +133,14 @@ def test_anthropic_invoke_with_message_history_and_system_instruction(
     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_with(
         model="claude-3-opus-20240229",
         system=system_instruction,
         messages=message_history,
         **model_params,
     )
 
-    assert llm.client.messages.create.call_count == 1  # type: ignore
+    assert llm.client.messages.create.call_count == 1
 
 
 def test_anthropic_invoke_with_message_history_validation_error(
@@ -178,7 +178,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(  # type: ignore
+    llm.async_client.messages.create.assert_awaited_once_with(
         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(  # type: ignore[attr-defined]
+    llm.client.chat.complete.assert_called_once_with(
         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(  # type: ignore[attr-defined]
+    llm.client.chat.complete.assert_called_once_with(
         messages=messages,
         model=model,
     )
 
-    assert llm.client.chat.complete.call_count == 1  # type: ignore
+    assert llm.client.chat.complete.call_count == 1
 
 
 @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(  # type: ignore[attr-defined]
+    llm.client.chat.assert_called_once_with(
         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(  # type: ignore[attr-defined]
+    llm.client.chat.assert_called_once_with(
         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(  # type: ignore[attr-defined]
+    llm.client.chat.assert_called_once_with(
         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(  # type: ignore[attr-defined]
+    llm.client.chat.assert_called_once_with(
         model=model, messages=messages, options=model_params
     )
-    assert llm.client.chat.call_count == 1  # type: ignore
+    assert llm.client.chat.call_count == 1
 
 
 @patch("builtins.__import__")