Added message history classes

Author: alexthomas93

src/neo4j_graphrag/message_history.py

@@ -0,0 +1,145 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional, Union
+
+import neo4j
+from pydantic import PositiveInt
+
+from neo4j_graphrag.llm.types import (
+    LLMMessage,
+)
+from neo4j_graphrag.types import (
+    Neo4jDriverModel,
+    Neo4jMessageHistoryModel,
+)
+
+CREATE_SESSION_NODE_QUERY = "MERGE (s:`{node_label}` {{id:$session_id}})"
+
+CLEAR_SESSION_QUERY = (
+    "MATCH (s:`{node_label}`)-[:LAST_MESSAGE]->(last_message) "
+    "WHERE s.id = $session_id MATCH p=(last_message)<-[:NEXT]-() "
+    "WITH p, length(p) AS length ORDER BY length DESC LIMIT 1 "
+    "UNWIND nodes(p) as node DETACH DELETE node;"
+)
+
+GET_MESSAGES_QUERY = (
+    "MATCH (s:`{node_label}`)-[:LAST_MESSAGE]->(last_message) "
+    "WHERE s.id = $session_id MATCH p=(last_message)<-[:NEXT*0.."
+    "{window}]-() WITH p, length(p) AS length "
+    "ORDER BY length DESC LIMIT 1 UNWIND reverse(nodes(p)) AS node "
+    "RETURN {{data:{{content: node.content}}, role:node.role}} AS result"
+)
+
+ADD_MESSAGE_QUERY = (
+    "MATCH (s:`{node_label}`) WHERE s.id = $session_id "
+    "OPTIONAL MATCH (s)-[lm:LAST_MESSAGE]->(last_message) "
+    "CREATE (s)-[:LAST_MESSAGE]->(new:Message) "
+    "SET new += {{role:$role, content:$content}} "
+    "WITH new, lm, last_message WHERE last_message IS NOT NULL "
+    "CREATE (last_message)-[:NEXT]->(new) "
+    "DELETE lm"
+)
+
+
+class MessageHistory(ABC):
+    @property
+    @abstractmethod
+    def messages(self) -> List[LLMMessage]: ...
+
+    @abstractmethod
+    def add_message(self, message: LLMMessage) -> None: ...
+
+    def add_messages(self, messages: List[LLMMessage]) -> None:
+        for message in messages:
+            self.add_message(message)
+
+    @abstractmethod
+    def clear(self) -> None: ...
+
+
+class InMemoryMessageHistory(MessageHistory):
+    def __init__(self, messages: List[LLMMessage] = []) -> None:
+        self._messages = messages
+
+    @property
+    def messages(self) -> List[LLMMessage]:
+        return self._messages
+
+    def add_message(self, message: LLMMessage) -> None:
+        self._messages.append(message)
+
+    def add_messages(self, messages: List[LLMMessage]) -> None:
+        self._messages.extend(messages)
+
+    def clear(self) -> None:
+        self._messages = []
+
+
+class Neo4jMessageHistory(MessageHistory):
+    def __init__(
+        self,
+        session_id: Union[str, int],
+        driver: neo4j.Driver,
+        node_label: str = "Session",
+        window: Optional[PositiveInt] = None,
+    ) -> None:
+        validated_data = Neo4jMessageHistoryModel(
+            session_id=session_id,
+            driver_model=Neo4jDriverModel(driver=driver),
+            node_label=node_label,
+            window=window,
+        )
+        self._driver = validated_data.driver_model.driver
+        self._session_id = validated_data.session_id
+        self._node_label = validated_data.node_label
+        self._window = (
+            "" if validated_data.window is None else validated_data.window - 1
+        )
+        # Create session node
+        self._driver.execute_query(
+            query_=CREATE_SESSION_NODE_QUERY.format(node_label=self._node_label),
+            parameters_={"session_id": self._session_id},
+        )
+
+    @property
+    def messages(self) -> List[LLMMessage]:
+        result = self._driver.execute_query(
+            query_=GET_MESSAGES_QUERY.format(
+                node_label=self._node_label, window=self._window
+            ),
+            parameters_={"session_id": self._session_id},
+        )
+        messages = [
+            LLMMessage(
+                content=el["result"]["data"]["content"],
+                role=el["result"]["role"],
+            )
+            for el in result.records
+        ]
+        return messages
+
+    @messages.setter
+    def messages(self, messages: List[LLMMessage]) -> None:
+        raise NotImplementedError(
+            "Direct assignment to 'messages' is not allowed."
+            " Use the 'add_messages' instead."
+        )
+
+    def add_message(self, message: LLMMessage) -> None:
+        self._driver.execute_query(
+            query_=ADD_MESSAGE_QUERY.format(node_label=self._node_label),
+            parameters_={
+                "role": message["role"],
+                "content": message["content"],
+                "session_id": self._session_id,
+            },
+        )
+
+    def clear(self) -> None:
+        self._driver.execute_query(
+            query_=CLEAR_SESSION_QUERY.format(node_label=self._node_label),
+            parameters_={"session_id": self._session_id},
+        )
+
+    def __del__(self) -> None:
+        if self._driver:
+            self._driver.close()

src/neo4j_graphrag/types.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import Any, Callable, Literal, Optional
+from typing import Any, Callable, Literal, Optional, Union
 
 import neo4j
 from pydantic import (
@@ -251,3 +251,10 @@ class Text2CypherRetrieverModel(BaseModel):
     result_formatter: Optional[Callable[[neo4j.Record], RetrieverResultItem]] = None
     custom_prompt: Optional[str] = None
     neo4j_database: Optional[str] = None
+
+
+class Neo4jMessageHistoryModel(BaseModel):
+    session_id: Union[str, int]
+    driver_model: Neo4jDriverModel
+    node_label: str = "Session"
+    window: Optional[PositiveInt] = None

tests/e2e/test_message_history_e2e.py

@@ -0,0 +1,88 @@
+import neo4j
+from neo4j_graphrag.llm.types import LLMMessage
+from neo4j_graphrag.message_history import Neo4jMessageHistory
+
+
+def test_neo4j_message_history_add_message(driver: neo4j.Driver) -> None:
+    driver.execute_query(query_="MATCH (n) DETACH DELETE n;")
+    message_history = Neo4jMessageHistory(session_id="123", driver=driver)
+    message_history.add_message(
+        LLMMessage(role="user", content="Hello"),
+    )
+    assert len(message_history.messages) == 1
+    assert message_history.messages[0]["role"] == "user"
+    assert message_history.messages[0]["content"] == "Hello"
+
+
+def test_neo4j_message_history_add_messages(driver: neo4j.Driver) -> None:
+    driver.execute_query(query_="MATCH (n) DETACH DELETE n;")
+    message_history = Neo4jMessageHistory(session_id="123", driver=driver)
+    message_history.add_messages(
+        [
+            LLMMessage(role="system", content="You are a helpful assistant."),
+            LLMMessage(role="user", content="Hello"),
+            LLMMessage(
+                role="assistant",
+                content="Hello, how may I help you today?",
+            ),
+            LLMMessage(role="user", content="I'd like to buy a new car."),
+            LLMMessage(
+                role="assistant",
+                content="I'd be happy to help you find the perfect car.",
+            ),
+        ]
+    )
+    assert len(message_history.messages) == 5
+    assert message_history.messages[0]["role"] == "system"
+    assert message_history.messages[0]["content"] == "You are a helpful assistant."
+    assert message_history.messages[1]["role"] == "user"
+    assert message_history.messages[1]["content"] == "Hello"
+    assert message_history.messages[2]["role"] == "assistant"
+    assert message_history.messages[2]["content"] == "Hello, how may I help you today?"
+    assert message_history.messages[3]["role"] == "user"
+    assert message_history.messages[3]["content"] == "I'd like to buy a new car."
+    assert message_history.messages[4]["role"] == "assistant"
+    assert (
+        message_history.messages[4]["content"]
+        == "I'd be happy to help you find the perfect car."
+    )
+
+
+def test_neo4j_message_history_clear(driver: neo4j.Driver) -> None:
+    driver.execute_query(query_="MATCH (n) DETACH DELETE n;")
+    message_history = Neo4jMessageHistory(session_id="123", driver=driver)
+    message_history.add_messages(
+        [
+            LLMMessage(role="system", content="You are a helpful assistant."),
+            LLMMessage(role="user", content="Hello"),
+        ]
+    )
+    assert len(message_history.messages) == 2
+    message_history.clear()
+    assert len(message_history.messages) == 0
+
+
+def test_neo4j_message_window_size(driver: neo4j.Driver) -> None:
+    driver.execute_query(query_="MATCH (n) DETACH DELETE n;")
+    message_history = Neo4jMessageHistory(session_id="123", driver=driver, window=1)
+    message_history.add_messages(
+        [
+            LLMMessage(role="system", content="You are a helpful assistant."),
+            LLMMessage(role="user", content="Hello"),
+            LLMMessage(
+                role="assistant",
+                content="Hello, how may I help you today?",
+            ),
+            LLMMessage(role="user", content="I'd like to buy a new car."),
+            LLMMessage(
+                role="assistant",
+                content="I'd be happy to help you find the perfect car.",
+            ),
+        ]
+    )
+    assert len(message_history.messages) == 1
+    assert (
+        message_history.messages[0]["content"]
+        == "I'd be happy to help you find the perfect car."
+    )
+    assert message_history.messages[0]["role"] == "assistant"

tests/unit/test_message_history.py

@@ -0,0 +1,93 @@
+from unittest.mock import MagicMock
+
+import pytest
+from neo4j_graphrag.llm.types import LLMMessage
+from neo4j_graphrag.message_history import InMemoryMessageHistory, Neo4jMessageHistory
+from pydantic import ValidationError
+
+
+def test_in_memory_message_history_add_message() -> None:
+    message_history = InMemoryMessageHistory()
+    message_history.add_message(
+        LLMMessage(role="user", content="may thy knife chip and shatter")
+    )
+    assert len(message_history.messages) == 1
+    assert message_history.messages[0]["role"] == "user"
+    assert message_history.messages[0]["content"] == "may thy knife chip and shatter"
+
+
+def test_in_memory_message_history_add_messages() -> None:
+    message_history = InMemoryMessageHistory()
+    message_history.add_messages(
+        [
+            LLMMessage(role="user", content="may thy knife chip and shatter"),
+            LLMMessage(
+                role="assistant",
+                content="He who controls the spice controls the universe.",
+            ),
+        ]
+    )
+    assert len(message_history.messages) == 2
+    assert message_history.messages[0]["role"] == "user"
+    assert message_history.messages[0]["content"] == "may thy knife chip and shatter"
+    assert message_history.messages[1]["role"] == "assistant"
+    assert (
+        message_history.messages[1]["content"]
+        == "He who controls the spice controls the universe."
+    )
+
+
+def test_in_memory_message_history_clear() -> None:
+    message_history = InMemoryMessageHistory()
+    message_history.add_messages(
+        [
+            LLMMessage(role="user", content="may thy knife chip and shatter"),
+            LLMMessage(
+                role="assistant",
+                content="He who controls the spice controls the universe.",
+            ),
+        ]
+    )
+    assert len(message_history.messages) == 2
+    message_history.clear()
+    assert len(message_history.messages) == 0
+
+
+def test_neo4j_message_history_invalid_session_id(driver: MagicMock) -> None:
+    with pytest.raises(ValidationError) as exc_info:
+        Neo4jMessageHistory(session_id=1.5, driver=driver, node_label="123", window=1)  # type: ignore[arg-type]
+    assert "Input should be a valid string" in str(exc_info.value)
+
+
+def test_neo4j_message_history_invalid_driver() -> None:
+    with pytest.raises(ValidationError) as exc_info:
+        Neo4jMessageHistory(session_id="123", driver=1.5, node_label="123", window=1)  # type: ignore[arg-type]
+    assert "Input should be a valid dictionary or instance of Neo4jDriver" in str(
+        exc_info.value
+    )
+
+
+def test_neo4j_message_history_invalid_node_label(driver: MagicMock) -> None:
+    with pytest.raises(ValidationError) as exc_info:
+        Neo4jMessageHistory(session_id="123", driver=driver, node_label=1.5, window=1)  # type: ignore[arg-type]
+    assert "Input should be a valid string" in str(exc_info.value)
+
+
+def test_neo4j_message_history_invalid_window(driver: MagicMock) -> None:
+    with pytest.raises(ValidationError) as exc_info:
+        Neo4jMessageHistory(
+            session_id="123", driver=driver, node_label="123", window=-1
+        )
+    assert "Input should be greater than 0" in str(exc_info.value)
+
+
+def test_neo4j_message_history_messages_setter(neo4j_driver: MagicMock) -> None:
+    message_history = Neo4jMessageHistory(session_id="123", driver=neo4j_driver)
+    with pytest.raises(NotImplementedError) as exc_info:
+        message_history.messages = [
+            LLMMessage(role="user", content="may thy knife chip and shatter"),
+        ]
+    assert (
+        str(exc_info.value)
+        == "Direct assignment to 'messages' is not allowed. Use the 'add_messages' instead."
+    )