Updated LLMMessage deprecation warning

Author: alexthomas93

src/neo4j_graphrag/llm/types.py

@@ -1,16 +1,20 @@
 import warnings
-from typing import Literal
+from typing import Any, Literal
 
 from pydantic import BaseModel
 
 from neo4j_graphrag.types import LLMMessage as _LLMMessage
 
-warnings.warn(
-    "LLMMessage has been moved to neo4j_graphrag.types. Please update your imports.",
-    DeprecationWarning,
-    stacklevel=2,
-)
-LLMMessage = _LLMMessage
+
+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 LLMResponse(BaseModel):