diff --git a/python/ai-server/src/ai_server/py_client/gaas/model.py b/python/ai-server/src/ai_server/py_client/gaas/model.py
index ecb7fea..64d3363 100644
--- a/python/ai-server/src/ai_server/py_client/gaas/model.py
+++ b/python/ai-server/src/ai_server/py_client/gaas/model.py
@@ -283,6 +283,8 @@ def get_chat_history(
                 history = self.model_engine.get_conversation_history()
                 messages = []
                 for msg in sorted(history, key=lambda x: x["DATE_CREATED"]):
+                    if "MESSAGE_DATA" not in msg:
+                        continue
                     if msg["MESSAGE_TYPE"] == "INPUT":
                         messages.append(HumanMessage(content=msg["MESSAGE_DATA"]))
                     elif msg["MESSAGE_TYPE"] == "RESPONSE":
@@ -333,7 +335,8 @@ def _create_chat_result(self, response: Dict[str, Any]) -> ChatResult:
                 # if this is a tool
                 # need to do a different return
 
-                if response["messageType"] == "TOOL":
+                # if response["messageType"] == "TOOL":
+                if response.get("messageType") == "TOOL":  # Returns None in case "messageType" is not present
                     tool_response = []
                     for m in message:
                         tool_response.append(