diff --git a/peak_assistant/streamlit/util/helpers.py b/peak_assistant/streamlit/util/helpers.py
index f25ddd4..cb94590 100644
--- a/peak_assistant/streamlit/util/helpers.py
+++ b/peak_assistant/streamlit/util/helpers.py
@@ -21,7 +21,8 @@
 # SPDX-License-Identifier: MIT
 
 from typing import List, Dict, Any, Optional, Tuple
-from autogen_agentchat.messages import TextMessage, UserMessage
+from autogen_agentchat.messages import TextMessage
+from autogen_core.models import UserMessage
 import streamlit as st
 import hashlib
 import html
diff --git a/tests/unit_tests/test_chat_history_conversion.py b/tests/unit_tests/test_chat_history_conversion.py
index 752aa34..77afeb9 100644
--- a/tests/unit_tests/test_chat_history_conversion.py
+++ b/tests/unit_tests/test_chat_history_conversion.py
@@ -1,6 +1,6 @@
 """Tests for convert_chat_history_to_user_messages and convert_chat_history_to_text_messages."""
-import pytest
-from autogen_agentchat.messages import TextMessage, UserMessage
+from autogen_agentchat.messages import TextMessage
+from autogen_core.models import UserMessage
 
 from peak_assistant.streamlit.util.helpers import (
     convert_chat_history_to_user_messages,