Updated more imports

alexthomas93 · Feb 21, 2025 · 685882b · 685882b
1 parent cf571b1
commit 685882b
Show file tree

Hide file tree

Showing 5 changed files with 6 additions and 7 deletions.
diff --git a/src/neo4j_graphrag/llm/anthropic_llm.py b/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

diff --git a/src/neo4j_graphrag/llm/cohere_llm.py b/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

diff --git a/src/neo4j_graphrag/llm/mistralai_llm.py b/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

diff --git a/src/neo4j_graphrag/llm/vertexai_llm.py b/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 (

diff --git a/src/neo4j_graphrag/message_history.py b/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,
 )