Support for conversations with message history (#234)

* Add system_instruction parameter

* Add chat_history parameter

* Add missing doc strings

* Open AI

* Add a summary of the chat history to the query embedding

* an extra llm call adds to latency and cost, but including the entire - or even part of - the chat history can potentially create a very large embedding context

* Anthropic

* Change return type of Anthropic get_messages()

* from list to anthropic.MessageParam

* Cohere

* upgrading from Cohere API v1 to API v2, as the v2 is handling chat history in a way that is consistent with the other providers

* Mistral

* VertexAI

* Formatting

* Fix mypy errors

* Ollama

* plus added the `options` parameter to the ollama `chat` call

* Override of the system message

* an idea of how to override the system instructions for some invokations

* Use TYPE_CHECKING for dev dependencies

* Formatting

* Rename `chat_history` to `message_history`

* Use BaseMessage class type

* for the type declaration of the `message_history` parameter

* System instruction override

* Revert BaseMessage class type

* bring back list[dicy[str,str]] type declaration for the `message_history` parameter

* Fix mypy errors

* Update tests

* Fix ollama NameError

* Fix NameError in unit tests

* at the same time as making mypy shut up

* Add TypeDict `LLMMessage`

* to help with the type declaration of the message history

* Simplify the retriever prompt

* Fix E2E tests

* Unit tests for the system instruction override

* Move and rename the prompts

* ... for query embedding and summarization to the GraphRAG class

* Update changelog

* Add missing parameter in example

* Add LLMMessage to the docs

* Update docs README

Author: leila-messallem

CHANGELOG.md

@@ -2,6 +2,14 @@
 
 ## Next
 
+### Added
+- Support for conversations with message history, including a new `message_history` parameter for LLM interactions.
+- Ability to include system instructions and override them for specific invocations.
+- Summarization of chat history to enhance query embedding and context handling.
+
+### Changed
+- Updated LLM implementations to handle message history consistently across providers.
+
 ## 1.3.0
 
 ### Added

docs/README.md

@@ -2,9 +2,11 @@
 
 Building the docs requires Python 3.8.1+
 
-Ensure the dev dependencies in `pyproject.toml` are installed.
+1. Ensure the dev dependencies in `pyproject.toml` are installed.
 
-From the root directory, run the Makefile:
+2. Add your changes to the appropriate `.rst` source file in `docs/source` directory.
+
+3. From the root directory, run the Makefile:
 
 ```
 make -C docs html

docs/source/types.rst

@@ -28,6 +28,12 @@ LLMResponse
 .. autoclass:: neo4j_graphrag.llm.types.LLMResponse
 
 
+LLMMessage
+===========
+
+.. autoclass:: neo4j_graphrag.llm.types.LLMMessage
+
+
 RagResultModel
 ==============
 

examples/customize/llms/custom_llm.py

@@ -1,21 +1,34 @@
 import random
 import string
-from typing import Any
+from typing import Any, Optional
 
 from neo4j_graphrag.llm import LLMInterface, LLMResponse
+from neo4j_graphrag.llm.types import LLMMessage
 
 
 class CustomLLM(LLMInterface):
-    def __init__(self, model_name: str, **kwargs: Any):
+    def __init__(
+        self, model_name: str, system_instruction: Optional[str] = None, **kwargs: Any
+    ):
         super().__init__(model_name, **kwargs)
 
-    def invoke(self, input: str) -> LLMResponse:
+    def invoke(
+        self,
+        input: str,
+        message_history: Optional[list[LLMMessage]] = None,
+        system_instruction: Optional[str] = None,
+    ) -> LLMResponse:
         content: str = (
             self.model_name + ": " + "".join(random.choices(string.ascii_letters, k=30))
         )
         return LLMResponse(content=content)
 
-    async def ainvoke(self, input: str) -> LLMResponse:
+    async def ainvoke(
+        self,
+        input: str,
+        message_history: Optional[list[LLMMessage]] = None,
+        system_instruction: Optional[str] = None,
+    ) -> LLMResponse:
         raise NotImplementedError()
 
 

src/neo4j_graphrag/generation/graphrag.py

@@ -27,6 +27,7 @@
 from neo4j_graphrag.generation.prompts import RagTemplate
 from neo4j_graphrag.generation.types import RagInitModel, RagResultModel, RagSearchModel
 from neo4j_graphrag.llm import LLMInterface
+from neo4j_graphrag.llm.types import LLMMessage
 from neo4j_graphrag.retrievers.base import Retriever
 from neo4j_graphrag.types import RetrieverResult
 
@@ -83,6 +84,7 @@ def __init__(
     def search(
         self,
         query_text: str = "",
+        message_history: Optional[list[LLMMessage]] = None,
         examples: str = "",
         retriever_config: Optional[dict[str, Any]] = None,
         return_context: bool | None = None,
@@ -99,14 +101,15 @@ def search(
 
 
         Args:
-            query_text (str): The user question
+            query_text (str): The user question.
+            message_history (Optional[list]): A collection previous messages, with each message having a specific role assigned.
             examples (str): Examples added to the LLM prompt.
-            retriever_config (Optional[dict]): Parameters passed to the retriever
+            retriever_config (Optional[dict]): Parameters passed to the retriever.
                 search method; e.g.: top_k
-            return_context (bool): Whether to append the retriever result to the final result (default: False)
+            return_context (bool): Whether to append the retriever result to the final result (default: False).
 
         Returns:
-            RagResultModel: The LLM-generated answer
+            RagResultModel: The LLM-generated answer.
 
         """
         if return_context is None:
@@ -124,18 +127,54 @@ def search(
             )
         except ValidationError as e:
             raise SearchValidationError(e.errors())
-        query_text = validated_data.query_text
+        query = self.build_query(validated_data.query_text, message_history)
         retriever_result: RetrieverResult = self.retriever.search(
-            query_text=query_text, **validated_data.retriever_config
+            query_text=query, **validated_data.retriever_config
         )
         context = "\n".join(item.content for item in retriever_result.items)
         prompt = self.prompt_template.format(
             query_text=query_text, context=context, examples=validated_data.examples
         )
         logger.debug(f"RAG: retriever_result={retriever_result}")
         logger.debug(f"RAG: prompt={prompt}")
-        answer = self.llm.invoke(prompt)
+        answer = self.llm.invoke(prompt, message_history)
         result: dict[str, Any] = {"answer": answer.content}
         if return_context:
             result["retriever_result"] = retriever_result
         return RagResultModel(**result)
+
+    def build_query(
+        self, query_text: str, message_history: Optional[list[LLMMessage]] = None
+    ) -> str:
+        summary_system_message = "You are a summarization assistant. Summarize the given text in no more than 300 words."
+        if message_history:
+            summarization_prompt = self.chat_summary_prompt(
+                message_history=message_history
+            )
+            summary = self.llm.invoke(
+                input=summarization_prompt,
+                system_instruction=summary_system_message,
+            ).content
+            return self.conversation_prompt(summary=summary, current_query=query_text)
+        return query_text
+
+    def chat_summary_prompt(self, message_history: list[LLMMessage]) -> str:
+        message_list = [
+            ": ".join([f"{value}" for _, value in message.items()])
+            for message in message_history
+        ]
+        history = "\n".join(message_list)
+        return f"""
+Summarize the message history:
+
+{history}
+"""
+
+    def conversation_prompt(self, summary: str, current_query: str) -> str:
+        return f"""
+Message Summary: 
+{summary}
+
+Current Query: 
+{current_query}
+"""

src/neo4j_graphrag/llm/anthropic_llm.py

@@ -13,11 +13,22 @@
 #  limitations under the License.
 from __future__ import annotations
 
-from typing import Any, Optional
+from typing import Any, Iterable, Optional, TYPE_CHECKING, cast
+
+from pydantic import ValidationError
 
 from neo4j_graphrag.exceptions import LLMGenerationError
 from neo4j_graphrag.llm.base import LLMInterface
-from neo4j_graphrag.llm.types import LLMResponse
+from neo4j_graphrag.llm.types import (
+    BaseMessage,
+    LLMMessage,
+    LLMResponse,
+    MessageList,
+    UserMessage,
+)
+
+if TYPE_CHECKING:
+    from anthropic.types.message_param import MessageParam
 
 
 class AnthropicLLM(LLMInterface):
@@ -26,6 +37,7 @@ class AnthropicLLM(LLMInterface):
     Args:
         model_name (str, optional): Name of the LLM to use. Defaults to "gemini-1.5-flash-001".
         model_params (Optional[dict], optional): Additional parameters passed to the model when text is sent to it. Defaults to None.
+        system_instruction: Optional[str], optional): Additional instructions for setting the behavior and context for the model in a conversation. Defaults to None.
         **kwargs (Any): Arguments passed to the model when for the class is initialised. Defaults to None.
 
     Raises:
@@ -49,6 +61,7 @@ def __init__(
         self,
         model_name: str,
         model_params: Optional[dict[str, Any]] = None,
+        system_instruction: Optional[str] = None,
         **kwargs: Any,
     ):
         try:
@@ -58,55 +71,86 @@ def __init__(
                 """Could not import Anthropic Python client.
                 Please install it with `pip install "neo4j-graphrag[anthropic]"`."""
             )
-        super().__init__(model_name, model_params)
+        super().__init__(model_name, model_params, system_instruction)
         self.anthropic = anthropic
         self.client = anthropic.Anthropic(**kwargs)
         self.async_client = anthropic.AsyncAnthropic(**kwargs)
 
-    def invoke(self, input: str) -> LLMResponse:
+    def get_messages(
+        self, input: str, message_history: Optional[list[LLMMessage]] = None
+    ) -> Iterable[MessageParam]:
+        messages: list[dict[str, str]] = []
+        if message_history:
+            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 invoke(
+        self,
+        input: str,
+        message_history: Optional[list[LLMMessage]] = None,
+        system_instruction: Optional[str] = None,
+    ) -> LLMResponse:
         """Sends text to the LLM and returns a response.
 
         Args:
             input (str): The text to send to the LLM.
+            message_history (Optional[list]): A collection previous messages, with each message having a specific role assigned.
+            system_instruction (Optional[str]): An option to override the llm system message for this invokation.
 
         Returns:
             LLMResponse: The response from the LLM.
         """
         try:
+            messages = self.get_messages(input, message_history)
+            system_message = (
+                system_instruction
+                if system_instruction is not None
+                else self.system_instruction
+            )
             response = self.client.messages.create(
                 model=self.model_name,
-                messages=[
-                    {
-                        "role": "user",
-                        "content": input,
-                    }
-                ],
+                system=system_message,  # type: ignore
+                messages=messages,
                 **self.model_params,
             )
-            return LLMResponse(content=response.content)
+            return LLMResponse(content=response.content)  # type: ignore
         except self.anthropic.APIError as e:
             raise LLMGenerationError(e)
 
-    async def ainvoke(self, input: str) -> LLMResponse:
+    async def ainvoke(
+        self,
+        input: str,
+        message_history: Optional[list[LLMMessage]] = None,
+        system_instruction: Optional[str] = None,
+    ) -> LLMResponse:
         """Asynchronously sends text to the LLM and returns a response.
 
         Args:
             input (str): The text to send to the LLM.
+            message_history (Optional[list]): A collection previous messages, with each message having a specific role assigned.
+            system_instruction (Optional[str]): An option to override the llm system message for this invokation.
 
         Returns:
             LLMResponse: The response from the LLM.
         """
         try:
+            messages = self.get_messages(input, message_history)
+            system_message = (
+                system_instruction
+                if system_instruction is not None
+                else self.system_instruction
+            )
             response = await self.async_client.messages.create(
                 model=self.model_name,
-                messages=[
-                    {
-                        "role": "user",
-                        "content": input,
-                    }
-                ],
+                system=system_message,  # type: ignore
+                messages=messages,
                 **self.model_params,
             )
-            return LLMResponse(content=response.content)
+            return LLMResponse(content=response.content)  # type: ignore
         except self.anthropic.APIError as e:
             raise LLMGenerationError(e)

src/neo4j_graphrag/llm/base.py

@@ -17,7 +17,7 @@
 from abc import ABC, abstractmethod
 from typing import Any, Optional
 
-from .types import LLMResponse
+from .types import LLMMessage, LLMResponse
 
 
 class LLMInterface(ABC):
@@ -26,24 +26,34 @@ class LLMInterface(ABC):
     Args:
         model_name (str): The name of the language model.
         model_params (Optional[dict], optional): Additional parameters passed to the model when text is sent to it. Defaults to None.
+        system_instruction: Optional[str], optional): Additional instructions for setting the behavior and context for the model in a conversation. Defaults to None.
         **kwargs (Any): Arguments passed to the model when for the class is initialised. Defaults to None.
     """
 
     def __init__(
         self,
         model_name: str,
         model_params: Optional[dict[str, Any]] = None,
+        system_instruction: Optional[str] = None,
         **kwargs: Any,
     ):
         self.model_name = model_name
         self.model_params = model_params or {}
+        self.system_instruction = system_instruction
 
     @abstractmethod
-    def invoke(self, input: str) -> LLMResponse:
+    def invoke(
+        self,
+        input: str,
+        message_history: Optional[list[LLMMessage]] = None,
+        system_instruction: Optional[str] = None,
+    ) -> LLMResponse:
         """Sends a text input to the LLM and retrieves a response.
 
         Args:
-            input (str): Text sent to the LLM
+            input (str): Text sent to the LLM.
+            message_history (Optional[list]): A collection previous messages, with each message having a specific role assigned.
+            system_instruction (Optional[str]): An option to override the llm system message for this invokation.
 
         Returns:
             LLMResponse: The response from the LLM.
@@ -53,11 +63,18 @@ def invoke(self, input: str) -> LLMResponse:
         """
 
     @abstractmethod
-    async def ainvoke(self, input: str) -> LLMResponse:
+    async def ainvoke(
+        self,
+        input: str,
+        message_history: Optional[list[LLMMessage]] = None,
+        system_instruction: Optional[str] = None,
+    ) -> LLMResponse:
         """Asynchronously sends a text input to the LLM and retrieves a response.
 
         Args:
-            input (str): Text sent to the LLM
+            input (str): Text sent to the LLM.
+            message_history (Optional[list]): A collection previous messages, with each message having a specific role assigned.
+            system_instruction (Optional[str]): An option to override the llm system message for this invokation.
 
         Returns:
             LLMResponse: The response from the LLM.