diff --git a/openai-core/src/main/scala/io/cequence/openaiscala/JsonFormats.scala b/openai-core/src/main/scala/io/cequence/openaiscala/JsonFormats.scala
index 81c9f0d7..6e8e369b 100644
--- a/openai-core/src/main/scala/io/cequence/openaiscala/JsonFormats.scala
+++ b/openai-core/src/main/scala/io/cequence/openaiscala/JsonFormats.scala
@@ -224,14 +224,12 @@ object JsonFormats {
         }
 
       case ChatRole.Tool =>
-        json.asOpt[AssistantToolMessage] match {
-          case Some(assistantToolMessage) => assistantToolMessage
-          case None                       => json.as[ToolMessage]
-        }
+        json.as[ToolMessage]
 
       case ChatRole.Assistant =>
-        json.asOpt[AssistantToolMessage] match {
-          case Some(assistantToolMessage) => assistantToolMessage
+        // if contains tool_calls, then it is AssistantToolMessage
+        (json \ "tool_calls").asOpt[JsArray] match {
+          case Some(_) => json.as[AssistantToolMessage]
           case None =>
             json.asOpt[AssistantMessage] match {
               case Some(assistantMessage) => assistantMessage