Updated more imports

Author: alexthomas93

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

@@ -19,8 +19,9 @@
 
 from neo4j_graphrag.exceptions import LLMGenerationError
 from neo4j_graphrag.llm.base import LLMInterface
-from neo4j_graphrag.llm.types import BaseMessage, LLMMessage, LLMResponse, MessageList
+from neo4j_graphrag.llm.types import BaseMessage, LLMResponse, MessageList
 from neo4j_graphrag.message_history import MessageHistory
+from neo4j_graphrag.types import LLMMessage
 
 try:
     from vertexai.generative_models import (

src/neo4j_graphrag/message_history.py

@@ -19,10 +19,8 @@
 import neo4j
 from pydantic import PositiveInt
 
-from neo4j_graphrag.llm.types import (
-    LLMMessage,
-)
 from neo4j_graphrag.types import (
+    LLMMessage,
     Neo4jDriverModel,
     Neo4jMessageHistoryModel,
 )