diff --git a/autoload/airline/extensions/nvimlsp.vim b/autoload/airline/extensions/nvimlsp.vim
index c4aa328a..4ccb1e12 100644
--- a/autoload/airline/extensions/nvimlsp.vim
+++ b/autoload/airline/extensions/nvimlsp.vim
@@ -26,7 +26,10 @@ function! airline#extensions#nvimlsp#get(type) abort
 
   let symbol = is_err ? error_symbol : warning_symbol
 
-  if luaeval("pcall(require, 'vim.lsp.diagnostic')")
+  if luaeval("pcall(require, 'vim.diagnostic')")
+    let severity = a:type == 'Warning' ? 'Warn' : a:type
+    let num = len(v:lua.vim.diagnostic.get(0, { 'severity': severity }))
+  elseif luaeval("pcall(require, 'vim.lsp.diagnostic')")
     let num = v:lua.vim.lsp.diagnostic.get_count(0, a:type)
   else
     let num = v:lua.vim.lsp.util.buf_diagnostics_count(a:type)