Mypy / ruff

Author: stellasia

examples/customize/llms/custom_llm.py

@@ -1,14 +1,13 @@
 import random
 import string
-from typing import Any, Awaitable, Callable, List, Optional, TypeVar, Union
+from typing import Any, Awaitable, Callable, Optional, TypeVar
 
 from neo4j_graphrag.llm import LLMInterface, LLMResponse
 from neo4j_graphrag.llm.rate_limit import (
     RateLimitHandler,
     # rate_limit_handler,
     # async_rate_limit_handler,
 )
-from neo4j_graphrag.message_history import MessageHistory
 from neo4j_graphrag.types import LLMMessage
 
 
@@ -18,38 +17,26 @@ def __init__(
     ):
         super().__init__(model_name, **kwargs)
 
-    # Optional: Apply rate limit handling to synchronous invoke method
-    # @rate_limit_handler
-    def invoke(
+    def _invoke(
         self,
-        input: str,
-        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
-        system_instruction: Optional[str] = None,
+        input: list[LLMMessage],
     ) -> LLMResponse:
         content: str = (
             self.model_name + ": " + "".join(random.choices(string.ascii_letters, k=30))
         )
         return LLMResponse(content=content)
 
-    # Optional: Apply rate limit handling to asynchronous ainvoke method
-    # @async_rate_limit_handler
-    async def ainvoke(
+    async def _ainvoke(
         self,
-        input: str,
-        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
-        system_instruction: Optional[str] = None,
+        input: list[LLMMessage],
     ) -> LLMResponse:
         raise NotImplementedError()
 
 
-llm = CustomLLM(
-    ""
-)  # if rate_limit_handler and async_rate_limit_handler decorators are used, the default rate limit handler will be applied automatically (retry with exponential backoff)
+llm = CustomLLM("")
 res: LLMResponse = llm.invoke("text")
 print(res.content)
 
-# If rate_limit_handler and async_rate_limit_handler decorators are used and you want to use a custom rate limit handler
-# Type variables for function signatures used in rate limit handlers
 F = TypeVar("F", bound=Callable[..., Any])
 AF = TypeVar("AF", bound=Callable[..., Awaitable[Any]])
 

examples/customize/llms/ollama_llm.py

@@ -17,7 +17,6 @@
 ]
 
 
-
 llm = OllamaLLM(
     model_name="orca-mini:latest",
     # model_params={"options": {"temperature": 0}, "format": "json"},

src/neo4j_graphrag/llm/anthropic_llm.py

@@ -13,8 +13,7 @@
 #  limitations under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Iterable, Optional
-
+from typing import TYPE_CHECKING, Any, Iterable, Optional, Union
 
 from neo4j_graphrag.exceptions import LLMGenerationError
 from neo4j_graphrag.llm.base import LLMInterface
@@ -28,6 +27,7 @@
 
 if TYPE_CHECKING:
     from anthropic.types.message_param import MessageParam
+    from anthropic import NotGiven
 
 
 class AnthropicLLM(LLMInterface):
@@ -78,16 +78,16 @@ def __init__(
     def get_messages(
         self,
         input: list[LLMMessage],
-    ) -> tuple[str, Iterable[MessageParam]]:
+    ) -> tuple[Union[str, NotGiven], Iterable[MessageParam]]:
         messages: list[MessageParam] = []
-        system_instruction = self.anthropic.NOT_GIVEN
+        system_instruction: Union[str, NotGiven] = self.anthropic.NOT_GIVEN
         for i in input:
             if i["role"] == "system":
                 system_instruction = i["content"]
             else:
                 messages.append(
                     self.anthropic.types.MessageParam(
-                        role=i["role"],  # type: ignore
+                        role=i["role"],
                         content=i["content"],
                     )
                 )

src/neo4j_graphrag/llm/mistralai_llm.py

@@ -39,7 +39,7 @@
 except ImportError:
     Mistral = None  # type: ignore
     SDKError = None  # type: ignore
-    Messages = Any
+    Messages = None  # type: ignore
 
 
 class MistralAILLM(LLMInterface):

src/neo4j_graphrag/llm/ollama_llm.py

@@ -71,10 +71,7 @@ def get_messages(
         self,
         input: list[LLMMessage],
     ) -> Sequence[Message]:
-        return [
-            self.ollama.Message(**i)
-            for i in input
-        ]
+        return [self.ollama.Message(**i) for i in input]
 
     def _invoke(
         self,

src/neo4j_graphrag/llm/openai_llm.py

@@ -25,7 +25,7 @@
     Iterable,
     Sequence,
     Union,
-    cast,
+    cast, Type,
 )
 
 from neo4j_graphrag.message_history import MessageHistory
@@ -44,9 +44,11 @@
 if TYPE_CHECKING:
     from openai.types.chat import (
         ChatCompletionMessageParam,
-        ChatCompletionToolParam, ChatCompletionUserMessageParam,
-        ChatCompletionSystemMessageParam, ChatCompletionAssistantMessageParam,
-)
+        ChatCompletionToolParam,
+        ChatCompletionUserMessageParam,
+        ChatCompletionSystemMessageParam,
+        ChatCompletionAssistantMessageParam,
+    )
     from openai import OpenAI, AsyncOpenAI
     from .rate_limit import RateLimitHandler
 else:
@@ -93,7 +95,7 @@ def get_messages(
     ) -> Iterable[ChatCompletionMessageParam]:
         chat_messages = []
         for m in messages:
-            message_type: ChatCompletionMessageParam
+            message_type: Type[ChatCompletionMessageParam]
             if m["role"] == "system":
                 message_type = ChatCompletionSystemMessageParam
             elif m["role"] == "user":
@@ -104,7 +106,7 @@ def get_messages(
                 raise ValueError(f"Unknown message type: {m['role']}")
             chat_messages.append(
                 message_type(
-                    role=m["role"],
+                    role=m["role"],  # type: ignore
                     content=m["content"],
                 )
             )