1 files changed, 9 insertions, 9 deletions
diff --git a/Omni/Agent/Provider.hs b/Omni/Agent/Provider.hs
index fd6920d..1bb4f04 100644
--- a/Omni/Agent/Provider.hs
+++ b/Omni/Agent/Provider.hs
@@ -589,6 +589,11 @@ parseStreamChunk obj = do
             _ -> "Unknown error"
       Just (StreamError errMsg)
     _ -> do
+      let usageChunk = case KeyMap.lookup "usage" obj of
+            Just usageVal -> case Aeson.fromJSON usageVal of
+              Aeson.Success usage -> Just (StreamDone (ChatResult (Message Assistant "" Nothing Nothing) (Just usage)))
+              _ -> Nothing
+            _ -> Nothing
       case KeyMap.lookup "choices" obj of
         Just (Aeson.Array choices) | not (null choices) -> do
           case toList choices of
@@ -603,15 +608,10 @@ parseStreamChunk obj = do
                           | not (null tcs) ->
                               parseToolCallDelta (toList tcs)
                         _ -> Nothing
-                  contentChunk <|> toolCallChunk
-                _ -> Nothing
-            _ -> Nothing
-        _ -> do
-          case KeyMap.lookup "usage" obj of
-            Just usageVal -> case Aeson.fromJSON usageVal of
-              Aeson.Success usage -> Just (StreamDone (ChatResult (Message Assistant "" Nothing Nothing) (Just usage)))
-              _ -> Nothing
-            _ -> Nothing
+                  contentChunk <|> toolCallChunk <|> usageChunk
+                _ -> usageChunk
+            _ -> usageChunk
+        _ -> usageChunk
 
 parseToolCallDelta :: [Aeson.Value] -> Maybe StreamChunk
 parseToolCallDelta [] = Nothing