llminterfacev2 support added to ollama

Author: serengil

src/neo4j_graphrag/llm/ollama_llm.py

@@ -12,23 +12,36 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-from __future__ import annotations
 
+# built-in dependencies
+from __future__ import annotations
 import warnings
-from typing import TYPE_CHECKING, Any, Iterable, List, Optional, Sequence, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Iterable,
+    List,
+    Optional,
+    Sequence,
+    Union,
+    cast,
+    overload,
+)
 
+# 3rd-party dependencies
 from pydantic import ValidationError
 
+# project dependencies
 from neo4j_graphrag.exceptions import LLMGenerationError
 from neo4j_graphrag.message_history import MessageHistory
 from neo4j_graphrag.types import LLMMessage
-
-from .base import LLMInterface
 from neo4j_graphrag.utils.rate_limit import (
     RateLimitHandler,
     rate_limit_handler,
     async_rate_limit_handler,
 )
+
+from .base import LLMInterface, LLMInterfaceV2
 from .types import (
     BaseMessage,
     LLMResponse,
@@ -40,8 +53,12 @@
 if TYPE_CHECKING:
     from ollama import Message
 
+# pylint: disable=redefined-builtin, arguments-differ, raise-missing-from, no-else-return
+
+
+class OllamaLLM(LLMInterface, LLMInterfaceV2):
+    """LLM wrapper for Ollama models."""
 
-class OllamaLLM(LLMInterface):
     def __init__(
         self,
         model_name: str,
@@ -78,28 +95,66 @@ def __init__(
             )
             self.model_params = {"options": self.model_params}
 
-    def get_messages(
+    # overloads for LLMInterface and LLMInterfaceV2 methods
+    @overload
+    def invoke(
         self,
         input: str,
         message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
         system_instruction: Optional[str] = None,
-    ) -> Sequence[Message]:
-        messages = []
-        if system_instruction:
-            messages.append(SystemMessage(content=system_instruction).model_dump())
-        if message_history:
-            if isinstance(message_history, MessageHistory):
-                message_history = message_history.messages
-            try:
-                MessageList(messages=cast(list[BaseMessage], message_history))
-            except ValidationError as e:
-                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
+    ) -> LLMResponse: ...
 
-    @rate_limit_handler
+    @overload
     def invoke(
+        self,
+        input: List[LLMMessage],
+    ) -> LLMResponse: ...
+
+    @overload
+    async def ainvoke(
+        self,
+        input: str,
+        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
+        system_instruction: Optional[str] = None,
+    ) -> LLMResponse: ...
+
+    @overload
+    async def ainvoke(
+        self,
+        input: List[LLMMessage],
+    ) -> LLMResponse: ...
+
+    # switching logics to LLMInterface or LLMInterfaceV2
+    def invoke(
+        self,
+        input: Union[str, List[LLMMessage]],
+        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
+        system_instruction: Optional[str] = None,
+    ) -> LLMResponse:
+        if isinstance(input, str):
+            return self.__legacy_invoke(input, message_history, system_instruction)
+        elif isinstance(input, list):
+            return self.__brand_new_invoke(input)
+        else:
+            raise ValueError(f"Invalid input type for invoke method - {type(input)}")
+
+    async def ainvoke(
+        self,
+        input: Union[str, List[LLMMessage]],
+        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
+        system_instruction: Optional[str] = None,
+    ) -> LLMResponse:
+        if isinstance(input, str):
+            return await self.__legacy_ainvoke(
+                input, message_history, system_instruction
+            )
+        elif isinstance(input, list):
+            return await self.__brand_new_ainvoke(input)
+        else:
+            raise ValueError(f"Invalid input type for ainvoke method - {type(input)}")
+
+    @rate_limit_handler
+    def __legacy_invoke(
         self,
         input: str,
         message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
@@ -129,8 +184,31 @@ def invoke(
         except self.ollama.ResponseError as e:
             raise LLMGenerationError(e)
 
+    def __brand_new_invoke(
+        self,
+        input: List[LLMMessage],
+    ) -> LLMResponse:
+        """Sends text to the LLM and returns a response.
+
+        Args:
+            input (str): The text to send to the LLM.
+
+        Returns:
+            LLMResponse: The response from the LLM.
+        """
+        try:
+            response = self.client.chat(
+                model=self.model_name,
+                messages=self.get_brand_new_messages(input),
+                **self.model_params,
+            )
+            content = response.message.content or ""
+            return LLMResponse(content=content)
+        except self.ollama.ResponseError as e:
+            raise LLMGenerationError(e)
+
     @async_rate_limit_handler
-    async def ainvoke(
+    async def __legacy_ainvoke(
         self,
         input: str,
         message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
@@ -163,3 +241,59 @@ async def ainvoke(
             return LLMResponse(content=content)
         except self.ollama.ResponseError as e:
             raise LLMGenerationError(e)
+
+    async def __brand_new_ainvoke(
+        self,
+        input: List[LLMMessage],
+    ) -> LLMResponse:
+        """Asynchronously sends a text input to the OpenAI chat
+        completion model and returns the response's content.
+
+        Args:
+            input (str): Text sent to the LLM.
+
+        Returns:
+            LLMResponse: The response from OpenAI.
+
+        Raises:
+            LLMGenerationError: If anything goes wrong.
+        """
+        try:
+            response = await self.async_client.chat(
+                model=self.model_name,
+                messages=self.get_brand_new_messages(input),
+                options=self.model_params,
+            )
+            content = response.message.content or ""
+            return LLMResponse(content=content)
+        except self.ollama.ResponseError as e:
+            raise LLMGenerationError(e)
+
+    # subsdiary methods
+    def get_messages(
+        self,
+        input: str,
+        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
+        system_instruction: Optional[str] = None,
+    ) -> Sequence[Message]:
+        """Constructs the message list for the Ollama chat API."""
+        messages = []
+        if system_instruction:
+            messages.append(SystemMessage(content=system_instruction).model_dump())
+        if message_history:
+            if isinstance(message_history, MessageHistory):
+                message_history = message_history.messages
+            try:
+                MessageList(messages=cast(list[BaseMessage], message_history))
+            except ValidationError as e:
+                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
+
+    def get_brand_new_messages(
+        self,
+        input: list[LLMMessage],
+    ) -> Sequence[Message]:
+        """Constructs the message list for the Ollama chat API."""
+        return [self.ollama.Message(**i) for i in input]