diff --git a/cldk/analysis/python/treesitter/python_sitter.py b/cldk/analysis/python/treesitter/python_sitter.py
index ec3ba656..df5dcc3f 100644
--- a/cldk/analysis/python/treesitter/python_sitter.py
+++ b/cldk/analysis/python/treesitter/python_sitter.py
@@ -197,6 +197,7 @@ def get_all_classes(self, module: str) -> List[PyClass]:
                                    full_signature=class_full_signature,
                                    methods=methods,
                                    super_classes=super_classes,
+                                   class_name=klass_name,
                                    is_test_class=is_test_class))
         return classes
 
diff --git a/cldk/models/python/models.py b/cldk/models/python/models.py
index ab08b94c..1f398f8a 100644
--- a/cldk/models/python/models.py
+++ b/cldk/models/python/models.py
@@ -45,6 +45,7 @@ class PyClass(BaseModel):
     full_signature: str
     super_classes: List[str]
     is_test_class: bool
+    class_name: str = None
     methods: List[PyMethod]