Updated more imports

Author: alexthomas93

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/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,