Moves LLMMessage to avoid a cicular import with LLM classes

Author: alexthomas93

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/types.py

@@ -1,14 +1,19 @@
-from typing import Literal, TypedDict
+import warnings
+from typing import Literal
 
 from pydantic import BaseModel
 
+from neo4j_graphrag.types import LLMMessage as _LLMMessage
 
-class LLMResponse(BaseModel):
-    content: str
+warnings.warn(
+    "LLMMessage has been moved to neo4j_graphrag.types. Please update your imports.",
+    DeprecationWarning,
+    stacklevel=2,
+)
+LLMMessage = _LLMMessage
 
 
-class LLMMessage(TypedDict):
-    role: Literal["system", "user", "assistant"]
+class LLMResponse(BaseModel):
     content: str
 
 

src/neo4j_graphrag/message_history.py

@@ -82,8 +82,8 @@ class InMemoryMessageHistory(MessageHistory):
 
     .. code-block:: python
 
-        from neo4j_graphrag.llm.types import LLMMessage
         from neo4j_graphrag.message_history import InMemoryMessageHistory
+        from neo4j_graphrag.types import LLMMessage
 
         history = InMemoryMessageHistory()
 
@@ -125,8 +125,8 @@ class Neo4jMessageHistory(MessageHistory):
     .. code-block:: python
 
         import neo4j
-        from neo4j_graphrag.llm.types import LLMMessage
         from neo4j_graphrag.message_history import Neo4jMessageHistory
+        from neo4j_graphrag.types import LLMMessage
 
         driver = neo4j.GraphDatabase.driver(URI, auth=AUTH)
 

src/neo4j_graphrag/types.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import Any, Callable, Literal, Optional, Union
+from typing import Any, Callable, Literal, Optional, TypedDict, Union
 
 import neo4j
 from pydantic import (
@@ -263,3 +263,8 @@ def validate_session_id(cls, v: Union[str, int]) -> Union[str, int]:
         if isinstance(v, str) and len(v) == 0:
             raise ValueError("session_id cannot be empty")
         return v
+
+
+class LLMMessage(TypedDict):
+    role: Literal["system", "user", "assistant"]
+    content: str

tests/e2e/test_graphrag_e2e.py

@@ -21,10 +21,9 @@
 from neo4j_graphrag.generation.graphrag import GraphRAG
 from neo4j_graphrag.generation.types import RagResultModel
 from neo4j_graphrag.llm import LLMResponse
-from neo4j_graphrag.llm.types import LLMMessage
 from neo4j_graphrag.message_history import Neo4jMessageHistory
 from neo4j_graphrag.retrievers import VectorCypherRetriever
-from neo4j_graphrag.types import RetrieverResult, RetrieverResultItem
+from neo4j_graphrag.types import LLMMessage, RetrieverResult, RetrieverResultItem
 
 from tests.e2e.conftest import BiologyEmbedder
 from tests.e2e.utils import build_data_objects, populate_neo4j

tests/e2e/test_message_history_e2e.py

@@ -13,8 +13,8 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 import neo4j
-from neo4j_graphrag.llm.types import LLMMessage
 from neo4j_graphrag.message_history import Neo4jMessageHistory
+from neo4j_graphrag.types import LLMMessage
 
 
 def test_neo4j_message_history_add_message(driver: neo4j.Driver) -> None:

tests/unit/llm/test_vertexai_llm.py

@@ -19,8 +19,8 @@
 
 import pytest
 from neo4j_graphrag.exceptions import LLMGenerationError
-from neo4j_graphrag.llm.types import LLMMessage
 from neo4j_graphrag.llm.vertexai_llm import VertexAILLM
+from neo4j_graphrag.types import LLMMessage
 from vertexai.generative_models import Content, Part
 
 

tests/unit/test_graphrag.py

@@ -21,9 +21,8 @@
 from neo4j_graphrag.generation.prompts import RagTemplate
 from neo4j_graphrag.generation.types import RagResultModel
 from neo4j_graphrag.llm import LLMResponse
-from neo4j_graphrag.llm.types import LLMMessage
 from neo4j_graphrag.message_history import InMemoryMessageHistory
-from neo4j_graphrag.types import RetrieverResult, RetrieverResultItem
+from neo4j_graphrag.types import LLMMessage, RetrieverResult, RetrieverResultItem
 
 
 def test_graphrag_prompt_template() -> None:

tests/unit/test_message_history.py

@@ -15,8 +15,8 @@
 from unittest.mock import MagicMock
 
 import pytest
-from neo4j_graphrag.llm.types import LLMMessage
 from neo4j_graphrag.message_history import InMemoryMessageHistory, Neo4jMessageHistory
+from neo4j_graphrag.types import LLMMessage
 from pydantic import ValidationError