Fix types and format

Author: oskarhane

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/experimental/pipeline/pipeline.py

@@ -251,9 +251,6 @@ def get_neo4j_viz_graph(self, hide_unused_outputs: bool = True) -> NeoVizGraph:
                     color="#4C8BF5",  # Blue for component nodes
                     caption_align=CaptionAlignment.CENTER,
                     caption_size=12,
-                    pinned=False,
-                    x=0,
-                    y=0,
                 )
             )
             node_counter += 1
@@ -270,9 +267,6 @@ def get_neo4j_viz_graph(self, hide_unused_outputs: bool = True) -> NeoVizGraph:
                         color="#34A853",  # Green for output nodes
                         caption_align=CaptionAlignment.CENTER,
                         caption_size=10,
-                        pinned=False,
-                        x=0,
-                        y=0,
                     )
                 )
                 # Connect component to its output

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,

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,
@@ -66,16 +66,22 @@ def test_anthropic_invoke_with_message_history_happy_path(mock_anthropic: Mock)
         "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."},
-    ]
+    from neo4j_graphrag.message_history import InMemoryMessageHistory
+    from neo4j_graphrag.types import LLMMessage
+
+    message_history = InMemoryMessageHistory()
+    message_history.add_message(
+        LLMMessage(role="user", content="When does the sun come up in the summer?")
+    )
+    message_history.add_message(
+        LLMMessage(role="assistant", content="Usually around 6am.")
+    )
     question = "What about next season?"
 
-    response = llm.invoke(question, message_history)  # type: ignore
+    response = llm.invoke(question, message_history)
     assert response.content == "generated text"
-    message_history.append({"role": "user", "content": question})
-    llm.client.messages.create.assert_called_once_with(  # type: ignore[attr-defined]
+    message_history.add_message(LLMMessage(role="user", content=question))
+    llm.client.messages.create.assert_called_once_with(
         messages=message_history,
         model="claude-3-opus-20240229",
         system=anthropic.NOT_GIVEN,
@@ -101,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(  # 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(
@@ -123,24 +129,30 @@ def test_anthropic_invoke_with_message_history_and_system_instruction(
         "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."},
-    ]
+    from neo4j_graphrag.message_history import InMemoryMessageHistory
+    from neo4j_graphrag.types import LLMMessage
+
+    message_history = InMemoryMessageHistory()
+    message_history.add_message(
+        LLMMessage(role="user", content="When does the sun come up in the summer?")
+    )
+    message_history.add_message(
+        LLMMessage(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
+    response = llm.invoke(question, message_history, 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]
+    message_history.add_message(LLMMessage(role="user", content=question))
+    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 +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(  # 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__")