diff --git a/pyproject.toml b/pyproject.toml
index 6c3fc86..68273ae 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -170,5 +170,8 @@ exclude_lines = [
 ]
 
 [[tool.mypy.overrides]]
-module = "simplemma.*"
+module = [
+  "simplemma.*",
+  "litellm.*"
+]
 ignore_missing_imports = true
\ No newline at end of file
diff --git a/src/banks/extensions/generate.py b/src/banks/extensions/generate.py
index 2fab63f..21c3618 100644
--- a/src/banks/extensions/generate.py
+++ b/src/banks/extensions/generate.py
@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: 2023-present Massimiliano Pippi <mpippi@gmail.com>
 #
 # SPDX-License-Identifier: MIT
-from typing import ClassVar, cast
+from typing import cast
 
 from jinja2 import nodes
 from jinja2.ext import Extension
@@ -22,7 +22,7 @@ class GenerateExtension(Extension):
     """
 
     # a set of names that trigger the extension.
-    tags: ClassVar[set] = {"generate"}
+    tags = {"generate"}  # noqa
 
     def parse(self, parser):
         # We get the line number of the first token so that we can give
diff --git a/src/banks/extensions/inference_endpoint.py b/src/banks/extensions/inference_endpoint.py
index 5a2311a..99fee6f 100644
--- a/src/banks/extensions/inference_endpoint.py
+++ b/src/banks/extensions/inference_endpoint.py
@@ -3,7 +3,6 @@
 # SPDX-License-Identifier: MIT
 import html
 import os
-from typing import ClassVar
 
 import requests
 from jinja2 import nodes
@@ -23,7 +22,7 @@ class HFInferenceEndpointsExtension(Extension):
     """
 
     # a set of names that trigger the extension.
-    tags: ClassVar[set] = {"inference_endpoint"}
+    tags = {"inference_endpoint"}  # noqa
 
     def parse(self, parser):
         # We get the line number of the first token so that we can give