Fixes import error & adds delete_session_node option to Neo4jMessageHistory (#282)

* Moves LLMMessage to avoid a cicular import with LLM classes

* Updated more imports

* Updates docs

* Updated more imports

* Updated Neo4jMessageHistory to allow for optional session node deletion

* Updated LLMMessage deprecation warning

Author: alexthomas93

docs/source/types.rst

@@ -31,7 +31,7 @@ LLMResponse
 LLMMessage
 ===========
 
-.. autoclass:: neo4j_graphrag.llm.types.LLMMessage
+.. autoclass:: neo4j_graphrag.types.LLMMessage
 
 
 RagResultModel

examples/customize/llms/custom_llm.py

@@ -3,8 +3,8 @@
 from typing import Any, List, Optional, Union
 
 from neo4j_graphrag.llm import LLMInterface, LLMResponse
-from neo4j_graphrag.llm.types import LLMMessage
 from neo4j_graphrag.message_history import MessageHistory
+from neo4j_graphrag.types import LLMMessage
 
 
 class CustomLLM(LLMInterface):

src/neo4j_graphrag/generation/graphrag.py

@@ -27,10 +27,9 @@
 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.message_history import MessageHistory
 from neo4j_graphrag.retrievers.base import Retriever
-from neo4j_graphrag.types import RetrieverResult
+from neo4j_graphrag.types import LLMMessage, RetrieverResult
 
 logger = logging.getLogger(__name__)
 

src/neo4j_graphrag/llm/anthropic_llm.py

@@ -21,12 +21,12 @@
 from neo4j_graphrag.llm.base import LLMInterface
 from neo4j_graphrag.llm.types import (
     BaseMessage,
-    LLMMessage,
     LLMResponse,
     MessageList,
     UserMessage,
 )
 from neo4j_graphrag.message_history import MessageHistory
+from neo4j_graphrag.types import LLMMessage
 
 if TYPE_CHECKING:
     from anthropic.types.message_param import MessageParam

src/neo4j_graphrag/llm/base.py

@@ -18,11 +18,9 @@
 from typing import Any, List, Optional, Union
 
 from neo4j_graphrag.message_history import MessageHistory
+from neo4j_graphrag.types import LLMMessage
 
-from .types import (
-    LLMMessage,
-    LLMResponse,
-)
+from .types import LLMResponse
 
 
 class LLMInterface(ABC):

src/neo4j_graphrag/llm/cohere_llm.py

@@ -22,13 +22,13 @@
 from neo4j_graphrag.llm.base import LLMInterface
 from neo4j_graphrag.llm.types import (
     BaseMessage,
-    LLMMessage,
     LLMResponse,
     MessageList,
     SystemMessage,
     UserMessage,
 )
 from neo4j_graphrag.message_history import MessageHistory
+from neo4j_graphrag.types import LLMMessage
 
 if TYPE_CHECKING:
     from cohere import ChatMessages

src/neo4j_graphrag/llm/mistralai_llm.py

@@ -23,13 +23,13 @@
 from neo4j_graphrag.llm.base import LLMInterface
 from neo4j_graphrag.llm.types import (
     BaseMessage,
-    LLMMessage,
     LLMResponse,
     MessageList,
     SystemMessage,
     UserMessage,
 )
 from neo4j_graphrag.message_history import MessageHistory
+from neo4j_graphrag.types import LLMMessage
 
 try:
     from mistralai import Messages, Mistral

src/neo4j_graphrag/llm/ollama_llm.py

@@ -20,11 +20,11 @@
 
 from neo4j_graphrag.exceptions import LLMGenerationError
 from neo4j_graphrag.message_history import MessageHistory
+from neo4j_graphrag.types import LLMMessage
 
 from .base import LLMInterface
 from .types import (
     BaseMessage,
-    LLMMessage,
     LLMResponse,
     MessageList,
     SystemMessage,

src/neo4j_graphrag/llm/openai_llm.py

@@ -20,12 +20,12 @@
 from pydantic import ValidationError
 
 from neo4j_graphrag.message_history import MessageHistory
+from neo4j_graphrag.types import LLMMessage
 
 from ..exceptions import LLMGenerationError
 from .base import LLMInterface
 from .types import (
     BaseMessage,
-    LLMMessage,
     LLMResponse,
     MessageList,
     SystemMessage,

src/neo4j_graphrag/llm/types.py

@@ -1,14 +1,23 @@
-from typing import Literal, TypedDict
+import warnings
+from typing import Any, Literal
 
 from pydantic import BaseModel
 
+from neo4j_graphrag.types import LLMMessage as _LLMMessage
 
-class LLMResponse(BaseModel):
-    content: str
 
+def __getattr__(name: str) -> Any:
+    if name == "LLMMessage":
+        warnings.warn(
+            "LLMMessage has been moved to neo4j_graphrag.types. Please update your imports.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return _LLMMessage
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
-class LLMMessage(TypedDict):
-    role: Literal["system", "user", "assistant"]
+
+class LLMResponse(BaseModel):
     content: str