Formatting

Author: leila-messallem

src/neo4j_graphrag/generation/graphrag.py

@@ -14,7 +14,6 @@
 #  limitations under the License.
 from __future__ import annotations
 
-import json
 import logging
 import warnings
 from typing import Any, Optional
@@ -25,7 +24,11 @@
     RagInitializationError,
     SearchValidationError,
 )
-from neo4j_graphrag.generation.prompts import RagTemplate, ChatSummaryTemplate, ConversationTemplate
+from neo4j_graphrag.generation.prompts import (
+    RagTemplate,
+    ChatSummaryTemplate,
+    ConversationTemplate,
+)
 from neo4j_graphrag.generation.types import RagInitModel, RagResultModel, RagSearchModel
 from neo4j_graphrag.llm import LLMInterface
 from neo4j_graphrag.retrievers.base import Retriever
@@ -142,10 +145,14 @@ def search(
         if return_context:
             result["retriever_result"] = retriever_result
         return RagResultModel(**result)
-    
+
     def build_query(self, query_text: str, chat_history: list[dict[str, str]]) -> str:
         if chat_history:
-            summarization_prompt = ChatSummaryTemplate().format(chat_history=chat_history)
+            summarization_prompt = ChatSummaryTemplate().format(
+                chat_history=chat_history
+            )
             summary = self.llm.invoke(summarization_prompt).content
-            return ConversationTemplate().format(summary=summary, current_query=query_text)
+            return ConversationTemplate().format(
+                summary=summary, current_query=query_text
+            )
         return query_text

src/neo4j_graphrag/generation/prompts.py

@@ -192,7 +192,7 @@ def format(
         text: str = "",
     ) -> str:
         return super().format(text=text, schema=schema, examples=examples)
-    
+
 
 class ChatSummaryTemplate(PromptTemplate):
     DEFAULT_TEMPLATE = """
@@ -203,8 +203,11 @@ class ChatSummaryTemplate(PromptTemplate):
     EXPECTED_INPUTS = ["chat_history"]
 
     def format(self, chat_history: list[dict[str, str]]) -> str:
-        message_list = [': '.join([f"{value}" for _, value in message.items()]) for message in chat_history]
-        history = '\n'.join(message_list)
+        message_list = [
+            ": ".join([f"{value}" for _, value in message.items()])
+            for message in chat_history
+        ]
+        history = "\n".join(message_list)
         return super().format(chat_history=history)
 
 

src/neo4j_graphrag/llm/anthropic_llm.py

@@ -82,7 +82,9 @@ def get_messages(self, input: str, chat_history: list) -> Iterable[MessageParam]
         messages.append(UserMessage(content=input).model_dump())
         return messages
 
-    def invoke(self, input: str, chat_history: Optional[list[dict[str, str]]] = None) -> LLMResponse:
+    def invoke(
+        self, input: str, chat_history: Optional[list[dict[str, str]]] = None
+    ) -> LLMResponse:
         """Sends text to the LLM and returns a response.
 
         Args:
@@ -95,9 +97,9 @@ def invoke(self, input: str, chat_history: Optional[list[dict[str, str]]] = None
         try:
             messages = self.get_messages(input, chat_history)
             response = self.client.messages.create(
-                model = self.model_name,
-                system = self.system_instruction,
-                messages = messages,
+                model=self.model_name,
+                system=self.system_instruction,
+                messages=messages,
                 **self.model_params,
             )
             return LLMResponse(content=response.content)
@@ -119,9 +121,9 @@ async def ainvoke(
         try:
             messages = self.get_messages(input, chat_history)
             response = await self.async_client.messages.create(
-                model = self.model_name,
-                system = self.system_instruction,
-                messages = messages,
+                model=self.model_name,
+                system=self.system_instruction,
+                messages=messages,
                 **self.model_params,
             )
             return LLMResponse(content=response.content)

src/neo4j_graphrag/llm/base.py

@@ -42,7 +42,9 @@ def __init__(
         self.system_instruction = system_instruction
 
     @abstractmethod
-    def invoke(self, input: str, chat_history: Optional[list[dict[str, str]]] = None) -> LLMResponse:
+    def invoke(
+        self, input: str, chat_history: Optional[list[dict[str, str]]] = None
+    ) -> LLMResponse:
         """Sends a text input to the LLM and retrieves a response.
 
         Args:

src/neo4j_graphrag/llm/cohere_llm.py

@@ -14,12 +14,17 @@
 #  limitations under the License.
 from __future__ import annotations
 
-from typing import Any, Iterable, Optional
+from typing import Any, Optional
 from pydantic import ValidationError
 
 from neo4j_graphrag.exceptions import LLMGenerationError
 from neo4j_graphrag.llm.base import LLMInterface
-from neo4j_graphrag.llm.types import LLMResponse, MessageList, SystemMessage, UserMessage
+from neo4j_graphrag.llm.types import (
+    LLMResponse,
+    MessageList,
+    SystemMessage,
+    UserMessage,
+)
 
 try:
     import cohere
@@ -69,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: list) -> ChatMessages:  # type: ignore
         messages = []
         if self.system_instruction:
             messages.append(SystemMessage(content=self.system_instruction).model_dump())
@@ -82,7 +87,9 @@ def get_messages(self, input: str, chat_history: list) -> ChatMessages: # type:
         messages.append(UserMessage(content=input).model_dump())
         return messages
 
-    def invoke(self, input: str, chat_history: Optional[list[dict[str, str]]] = None) -> LLMResponse:
+    def invoke(
+        self, input: str, chat_history: Optional[list[dict[str, str]]] = None
+    ) -> LLMResponse:
         """Sends text to the LLM and returns a response.
 
         Args:

src/neo4j_graphrag/llm/mistralai_llm.py

@@ -15,12 +15,17 @@
 from __future__ import annotations
 
 import os
-from typing import Any, Optional, Union
+from typing import Any, Optional
 from pydantic import ValidationError
 
 from neo4j_graphrag.exceptions import LLMGenerationError
 from neo4j_graphrag.llm.base import LLMInterface
-from neo4j_graphrag.llm.types import LLMResponse, MessageList, SystemMessage, UserMessage
+from neo4j_graphrag.llm.types import (
+    LLMResponse,
+    MessageList,
+    SystemMessage,
+    UserMessage,
+)
 
 try:
     from mistralai import Mistral, Messages
@@ -58,7 +63,7 @@ def __init__(
         if api_key is None:
             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]:
         messages = []
         if self.system_instruction:
@@ -72,7 +77,9 @@ def get_messages(self, input: str, chat_history: list) -> list[Messages]:
         messages.append(UserMessage(content=input).model_dump())
         return messages
 
-    def invoke(self, input: str, chat_history: Optional[list[dict[str, str]]] = None) -> LLMResponse:
+    def invoke(
+        self, input: str, chat_history: Optional[list[dict[str, str]]] = None
+    ) -> LLMResponse:
         """Sends a text input to the Mistral chat completion model
         and returns the response's content.
 

src/neo4j_graphrag/llm/openai_llm.py

@@ -61,9 +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: list
     ) -> Iterable[ChatCompletionMessageParam]:
         messages = []
         if self.system_instruction:
@@ -77,7 +75,9 @@ def get_messages(
         messages.append(UserMessage(content=input).model_dump())
         return messages
 
-    def invoke(self, input: str, chat_history: Optional[list[dict[str, str]]] = None) -> LLMResponse:
+    def invoke(
+        self, input: str, chat_history: Optional[list[dict[str, str]]] = None
+    ) -> LLMResponse:
         """Sends a text input to the OpenAI chat completion model
         and returns the response's content.
 

src/neo4j_graphrag/llm/types.py

@@ -14,6 +14,7 @@ class BaseMessage(BaseModel):
 class UserMessage(BaseMessage):
     role: Literal["user"] = "user"
 
+
 class SystemMessage(BaseMessage):
     role: Literal["system"] = "system"
 

src/neo4j_graphrag/llm/vertexai_llm.py

@@ -76,25 +76,34 @@ def __init__(
             model_name=model_name, system_instruction=[system_instruction], **kwargs
         )
 
-
     def get_messages(self, input: str, chat_history: list[str]) -> list[Content]:
         messages = []
         if chat_history:
             try:
                 MessageList(messages=chat_history)
             except ValidationError as e:
                 raise LLMGenerationError(e.errors()) from e
-            
+
             for message in chat_history:
                 if message.get("role") == "user":
-                    messages.append(Content(role="user", parts=[Part.from_text(message.get("content"))]))
+                    messages.append(
+                        Content(
+                            role="user", parts=[Part.from_text(message.get("content"))]
+                        )
+                    )
                 elif message.get("role") == "assistant":
-                    messages.append(Content(role="model", parts=[Part.from_text(message.get("content"))]))
+                    messages.append(
+                        Content(
+                            role="model", parts=[Part.from_text(message.get("content"))]
+                        )
+                    )
 
         messages.append(Content(role="user", parts=[Part.from_text(input)]))
         return messages
 
-    def invoke(self, input: str, chat_history: Optional[list[dict[str, str]]] = None) -> LLMResponse:
+    def invoke(
+        self, input: str, chat_history: Optional[list[dict[str, str]]] = None
+    ) -> LLMResponse:
         """Sends text to the LLM and returns a response.
 
         Args:

tests/unit/llm/test_anthropic_llm.py

@@ -13,11 +13,8 @@
 #  limitations under the License.
 from __future__ import annotations
 
-import sys
-from typing import Generator
 from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
-import anthropic
 import pytest
 from neo4j_graphrag.exceptions import LLMGenerationError
 from neo4j_graphrag.llm.anthropic_llm import AnthropicLLM
@@ -54,17 +51,21 @@ def test_anthropic_invoke_with_chat_history_happy_path(mock_anthropic: Mock) ->
     )
     model_params = {"temperature": 0.3}
     system_instruction = "You are a helpful assistant."
-    llm = AnthropicLLM("claude-3-opus-20240229", model_params=model_params, system_instruction=system_instruction)
+    llm = AnthropicLLM(
+        "claude-3-opus-20240229",
+        model_params=model_params,
+        system_instruction=system_instruction,
+    )
     chat_history = [
         {"role": "user", "content": "When does the sun come up in the summer?"},
         {"role": "assistant", "content": "Usually around 6am."},
     ]
     question = "What about next season?"
-    
+
     response = llm.invoke(question, chat_history)
     assert response.content == "generated text"
     chat_history.append({"role": "user", "content": question})
-    llm.client.messages.create.assert_called_once_with(  # type: ignore
+    llm.client.messages.create.assert_called_once_with(
         messages=chat_history,
         model="claude-3-opus-20240229",
         system=system_instruction,
@@ -73,13 +74,19 @@ def test_anthropic_invoke_with_chat_history_happy_path(mock_anthropic: Mock) ->
 
 
 @patch("neo4j_graphrag.llm.anthropic_llm.anthropic.Anthropic")
-def test_anthropic_invoke_with_chat_history_validation_error(mock_anthropic: Mock) -> None:
+def test_anthropic_invoke_with_chat_history_validation_error(
+    mock_anthropic: Mock,
+) -> None:
     mock_anthropic.return_value.messages.create.return_value = MagicMock(
         content="generated text"
     )
     model_params = {"temperature": 0.3}
     system_instruction = "You are a helpful assistant."
-    llm = AnthropicLLM("claude-3-opus-20240229", model_params=model_params, system_instruction=system_instruction)
+    llm = AnthropicLLM(
+        "claude-3-opus-20240229",
+        model_params=model_params,
+        system_instruction=system_instruction,
+    )
     chat_history = [
         {"role": "human", "content": "When does the sun come up in the summer?"},
         {"role": "assistant", "content": "Usually around 6am."},