Fix mypy errors

Author: leila-messallem

src/neo4j_graphrag/generation/graphrag.py

@@ -146,7 +146,7 @@ def search(
             result["retriever_result"] = retriever_result
         return RagResultModel(**result)
 
-    def build_query(self, query_text: str, chat_history: list[dict[str, str]]) -> str:
+    def build_query(self, query_text: str, chat_history: Optional[list[dict[str, str]]] = None) -> str:
         if chat_history:
             summarization_prompt = ChatSummaryTemplate().format(
                 chat_history=chat_history

src/neo4j_graphrag/llm/anthropic_llm.py

@@ -71,7 +71,7 @@ def __init__(
         self.client = anthropic.Anthropic(**kwargs)
         self.async_client = anthropic.AsyncAnthropic(**kwargs)
 
-    def get_messages(self, input: str, chat_history: list) -> Iterable[MessageParam]:
+    def get_messages(self, input: str, chat_history: Optional[list[Any]] = None) -> Iterable[MessageParam]:
         messages = []
         if chat_history:
             try:
@@ -83,7 +83,7 @@ def get_messages(self, input: str, chat_history: list) -> Iterable[MessageParam]
         return messages
 
     def invoke(
-        self, input: str, chat_history: Optional[list[dict[str, str]]] = None
+        self, input: str, chat_history: Optional[list[Any]] = None
     ) -> LLMResponse:
         """Sends text to the LLM and returns a response.
 
@@ -107,7 +107,7 @@ def invoke(
             raise LLMGenerationError(e)
 
     async def ainvoke(
-        self, input: str, chat_history: Optional[list[dict[str, str]]] = None
+        self, input: str, chat_history: Optional[list[Any]] = None
     ) -> LLMResponse:
         """Asynchronously sends text to the LLM and returns a response.
 

src/neo4j_graphrag/llm/cohere_llm.py

@@ -74,7 +74,7 @@ def __init__(
         self.client = cohere.ClientV2(**kwargs)
         self.async_client = cohere.AsyncClientV2(**kwargs)
 
-    def get_messages(self, input: str, chat_history: list) -> ChatMessages:  # type: ignore
+    def get_messages(self, input: str, chat_history: Optional[list[Any]] = None) -> ChatMessages:
         messages = []
         if self.system_instruction:
             messages.append(SystemMessage(content=self.system_instruction).model_dump())
@@ -88,7 +88,7 @@ def get_messages(self, input: str, chat_history: list) -> ChatMessages:  # type:
         return messages
 
     def invoke(
-        self, input: str, chat_history: Optional[list[dict[str, str]]] = None
+        self, input: str, chat_history: Optional[list[Any]] = None
     ) -> LLMResponse:
         """Sends text to the LLM and returns a response.
 
@@ -112,7 +112,7 @@ def invoke(
         )
 
     async def ainvoke(
-        self, input: str, chat_history: Optional[list[dict[str, str]]] = None
+        self, input: str, chat_history: Optional[list[Any]] = None
     ) -> LLMResponse:
         """Asynchronously sends text to the LLM and returns a response.
 

src/neo4j_graphrag/llm/mistralai_llm.py

@@ -31,8 +31,8 @@
     from mistralai import Mistral, Messages
     from mistralai.models.sdkerror import SDKError
 except ImportError:
-    Mistral = None  # type: ignore
-    SDKError = None  # type: ignore
+    Mistral = None
+    SDKError = None
 
 
 class MistralAILLM(LLMInterface):
@@ -64,7 +64,7 @@ def __init__(
             api_key = os.getenv("MISTRAL_API_KEY", "")
         self.client = Mistral(api_key=api_key, **kwargs)
 
-    def get_messages(self, input: str, chat_history: list) -> list[Messages]:
+    def get_messages(self, input: str, chat_history: Optional[list[Any]] = None) -> list[Messages]:
         messages = []
         if self.system_instruction:
             messages.append(SystemMessage(content=self.system_instruction).model_dump())
@@ -78,7 +78,7 @@ def get_messages(self, input: str, chat_history: list) -> list[Messages]:
         return messages
 
     def invoke(
-        self, input: str, chat_history: Optional[list[dict[str, str]]] = None
+        self, input: str, chat_history: Optional[list[Any]] = None
     ) -> LLMResponse:
         """Sends a text input to the Mistral chat completion model
         and returns the response's content.
@@ -110,7 +110,7 @@ def invoke(
             raise LLMGenerationError(e)
 
     async def ainvoke(
-        self, input: str, chat_history: Optional[list[dict[str, str]]] = None
+        self, input: str, chat_history: Optional[list[Any]] = None
     ) -> LLMResponse:
         """Asynchronously sends a text input to the MistralAI chat
         completion model and returns the response's content.

src/neo4j_graphrag/llm/openai_llm.py

@@ -61,7 +61,7 @@ def __init__(
         super().__init__(model_name, model_params, system_instruction)
 
     def get_messages(
-        self, input: str, chat_history: list
+        self, input: str, chat_history: Optional[list[Any]] = None
     ) -> Iterable[ChatCompletionMessageParam]:
         messages = []
         if self.system_instruction:
@@ -76,7 +76,7 @@ def get_messages(
         return messages
 
     def invoke(
-        self, input: str, chat_history: Optional[list[dict[str, str]]] = None
+        self, input: str, chat_history: Optional[list[Any]] = None
     ) -> LLMResponse:
         """Sends a text input to the OpenAI chat completion model
         and returns the response's content.
@@ -103,7 +103,7 @@ def invoke(
             raise LLMGenerationError(e)
 
     async def ainvoke(
-        self, input: str, chat_history: Optional[list[dict[str, str]]] = None
+        self, input: str, chat_history: Optional[list[Any]] = None
     ) -> LLMResponse:
         """Asynchronously sends a text input to the OpenAI chat
         completion model and returns the response's content.

src/neo4j_graphrag/llm/types.py

@@ -7,7 +7,7 @@ class LLMResponse(BaseModel):
 
 
 class BaseMessage(BaseModel):
-    role: Literal["user", "assistant"]
+    role: Literal["user", "assistant", "system"]
     content: str
 
 

src/neo4j_graphrag/llm/vertexai_llm.py

@@ -76,7 +76,7 @@ def __init__(
             model_name=model_name, system_instruction=[system_instruction], **kwargs
         )
 
-    def get_messages(self, input: str, chat_history: list[str]) -> list[Content]:
+    def get_messages(self, input: str, chat_history: Optional[list[Any]] = None) -> list[Content]:
         messages = []
         if chat_history:
             try:
@@ -102,7 +102,7 @@ def get_messages(self, input: str, chat_history: list[str]) -> list[Content]:
         return messages
 
     def invoke(
-        self, input: str, chat_history: Optional[list[dict[str, str]]] = None
+        self, input: str, chat_history: Optional[list[Any]] = None
     ) -> LLMResponse:
         """Sends text to the LLM and returns a response.
 
@@ -121,7 +121,7 @@ def invoke(
             raise LLMGenerationError(e)
 
     async def ainvoke(
-        self, input: str, chat_history: Optional[list[dict[str, str]]] = None
+        self, input: str, chat_history: Optional[list[Any]] = None
     ) -> LLMResponse:
         """Asynchronously sends text to the LLM and returns a response.