diff --git a/src/class_diagram/generators/plantuml/class_diagram_generator.cc b/src/class_diagram/generators/plantuml/class_diagram_generator.cc
index 77269e6b..92b749fd 100644
--- a/src/class_diagram/generators/plantuml/class_diagram_generator.cc
+++ b/src/class_diagram/generators/plantuml/class_diagram_generator.cc
@@ -489,6 +489,8 @@ void generator::generate_relationships(
 
 void generator::generate(std::ostream &ostr) const
 {
+    update_context();
+
     ostr << "@startuml" << '\n';
 
     generate_plantuml_directives(ostr, m_config.puml().before);
diff --git a/src/class_diagram/model/class.h b/src/class_diagram/model/class.h
index a894ada6..c6e1cfd1 100644
--- a/src/class_diagram/model/class.h
+++ b/src/class_diagram/model/class.h
@@ -42,6 +42,8 @@ public:
     class_ &operator=(const class_ &) = delete;
     class_ &operator=(class_ &&) = delete;
 
+    std::string type_name() const override { return "class"; }
+
     bool is_struct() const;
     void is_struct(bool is_struct);
 
diff --git a/src/class_diagram/model/diagram.cc b/src/class_diagram/model/diagram.cc
index 1a6098aa..4c428e40 100644
--- a/src/class_diagram/model/diagram.cc
+++ b/src/class_diagram/model/diagram.cc
@@ -271,6 +271,28 @@ std::string diagram::to_alias(
     throw error::uml_alias_missing(fmt::format("Missing alias for {}", id));
 }
 
+inja::json diagram::context() const {
+    inja::json ctx;
+    ctx["name"] = name();
+    ctx["type"] = "class";
+
+    inja::json::array_t elements{};
+
+    // Add classes
+    for(const auto &c : classes()) {
+        elements.emplace_back(c.get().context());
+    }
+
+    // Add enums
+    for(const auto &e : enums()) {
+        elements.emplace_back(e.get().context());
+    }
+
+    ctx["elements"] = elements;
+
+    return ctx;
+}
+
 }
 
 namespace clanguml::common::model {
diff --git a/src/class_diagram/model/diagram.h b/src/class_diagram/model/diagram.h
index 0f3d033e..0fda26ce 100644
--- a/src/class_diagram/model/diagram.h
+++ b/src/class_diagram/model/diagram.h
@@ -87,6 +87,8 @@ public:
     bool has_element(
         const clanguml::common::model::diagram_element::id_t id) const override;
 
+    inja::json context() const override;
+
 private:
     common::reference_vector<class_> classes_;
 
diff --git a/src/class_diagram/model/enum.h b/src/class_diagram/model/enum.h
index e6d93a7b..4cbfc72a 100644
--- a/src/class_diagram/model/enum.h
+++ b/src/class_diagram/model/enum.h
@@ -34,6 +34,8 @@ public:
     enum_ &operator=(const enum_ &) = delete;
     enum_ &operator=(enum_ &&) = delete;
 
+    std::string type_name() const override { return "enum"; }
+
     // TODO: Do we need this?
     friend bool operator==(const enum_ &l, const enum_ &r);
 
diff --git a/src/common/generators/plantuml/generator.h b/src/common/generators/plantuml/generator.h
index ab0c73fd..366f44c1 100644
--- a/src/common/generators/plantuml/generator.h
+++ b/src/common/generators/plantuml/generator.h
@@ -69,6 +69,8 @@ public:
     template <typename E>
     void generate_link(std::ostream &ostr, const E &e) const;
 
+    void update_context() const;
+
 protected:
     const inja::json &context() const;
 
@@ -85,7 +87,7 @@ protected:
     ConfigType &m_config;
     DiagramType &m_model;
     mutable std::set<std::string> m_generated_aliases;
-    inja::json m_context;
+    mutable inja::json m_context;
     mutable inja::Environment m_env;
 };
 
@@ -372,9 +374,11 @@ template <typename C, typename D> void generator<C, D>::init_context()
         m_context["git"]["commit"] = m_config.git().commit;
         m_context["git"]["toplevel"] = m_config.git().toplevel;
     }
+}
 
-    m_context["diagram"]["name"] = m_config.name;
-    m_context["diagram"]["type"] = to_string(m_config.type());
+template <typename C, typename D> void generator<C, D>::update_context() const
+{
+    m_context["diagram"] = m_model.context();
 }
 
 template <typename C, typename D> void generator<C, D>::init_env()
diff --git a/src/common/model/diagram.h b/src/common/model/diagram.h
index 4483072d..0bca1643 100644
--- a/src/common/model/diagram.h
+++ b/src/common/model/diagram.h
@@ -74,6 +74,8 @@ public:
     virtual bool should_include(
         const namespace_ &ns, const std::string &name) const;
 
+    virtual inja::json context() const = 0;
+
 private:
     std::string name_;
     std::unique_ptr<diagram_filter> filter_;
diff --git a/src/common/model/diagram_element.h b/src/common/model/diagram_element.h
index 2d58bb62..6a6db750 100644
--- a/src/common/model/diagram_element.h
+++ b/src/common/model/diagram_element.h
@@ -48,6 +48,8 @@ public:
 
     std::string name() const { return name_; }
 
+    virtual std::string type_name() const { return "__undefined__"; };
+
     virtual std::string full_name(bool /*relative*/) const { return name(); }
 
     std::vector<relationship> &relationships();
diff --git a/src/common/model/element.cc b/src/common/model/element.cc
index dd789d92..61ef96fc 100644
--- a/src/common/model/element.cc
+++ b/src/common/model/element.cc
@@ -40,6 +40,7 @@ inja::json element::context() const
 {
     inja::json ctx;
     ctx["name"] = name();
+    ctx["type"] = type_name();
     ctx["alias"] = alias();
     ctx["full_name"] = full_name(false);
     ctx["namespace"] = get_namespace().to_string();
diff --git a/src/common/model/nested_trait.h b/src/common/model/nested_trait.h
index 285f69f2..60b6f9b1 100644
--- a/src/common/model/nested_trait.h
+++ b/src/common/model/nested_trait.h
@@ -71,7 +71,7 @@ public:
             return dynamic_cast<nested_trait<T, Path> &>(parent.value())
                 .template add_element<V>(std::move(p));
         else {
-            spdlog::error("No parent element found at: {}", path.to_string());
+            spdlog::info("No parent element found at: {}", path.to_string());
             throw std::runtime_error(
                 "No parent element found for " + path.to_string());
         }
diff --git a/src/common/model/package.h b/src/common/model/package.h
index 49a6c375..c688f94d 100644
--- a/src/common/model/package.h
+++ b/src/common/model/package.h
@@ -42,6 +42,8 @@ public:
     package &operator=(const package &) = delete;
     package &operator=(package &&) = delete;
 
+    std::string type_name() const override { return "package"; }
+
     std::string full_name(bool relative) const override;
 
     bool is_deprecated() const;
diff --git a/src/include_diagram/model/diagram.cc b/src/include_diagram/model/diagram.cc
index be0244be..15a69379 100644
--- a/src/include_diagram/model/diagram.cc
+++ b/src/include_diagram/model/diagram.cc
@@ -135,6 +135,22 @@ diagram::files() const
     return files_;
 }
 
+inja::json diagram::context() const {
+    inja::json ctx;
+    ctx["name"] = name();
+    ctx["type"] = "include";
+
+    inja::json::array_t elements{};
+
+    // Add files and directories
+    for(const auto &f : files()) {
+        elements.emplace_back(f.get().context());
+    }
+
+    ctx["elements"] = elements;
+
+    return ctx;
+}
 }
 
 namespace clanguml::common::model {
diff --git a/src/include_diagram/model/diagram.h b/src/include_diagram/model/diagram.h
index 11dda724..8be09603 100644
--- a/src/include_diagram/model/diagram.h
+++ b/src/include_diagram/model/diagram.h
@@ -59,6 +59,8 @@ public:
 
     const common::reference_vector<common::model::source_file> &files() const;
 
+    inja::json context() const override;
+
 private:
     common::reference_vector<common::model::source_file> files_;
 };
diff --git a/src/package_diagram/model/diagram.cc b/src/package_diagram/model/diagram.cc
index 3f633682..4b78296b 100644
--- a/src/package_diagram/model/diagram.cc
+++ b/src/package_diagram/model/diagram.cc
@@ -94,6 +94,22 @@ std::string diagram::to_alias(
 
     return {};
 }
+
+inja::json diagram::context() const {
+    inja::json ctx;
+    ctx["name"] = name();
+    ctx["type"] = "package";
+
+    inja::json::array_t elements{};
+
+    for(const auto &p : packages()) {
+        elements.emplace_back(p.get().context());
+    }
+
+    ctx["elements"] = elements;
+
+    return ctx;
+}
 }
 
 namespace clanguml::common::model {
diff --git a/src/package_diagram/model/diagram.h b/src/package_diagram/model/diagram.h
index dacb6b2c..72c762c7 100644
--- a/src/package_diagram/model/diagram.h
+++ b/src/package_diagram/model/diagram.h
@@ -59,6 +59,8 @@ public:
     std::string to_alias(
         const clanguml::common::model::diagram_element::id_t) const;
 
+    inja::json context() const override;
+
 private:
     common::reference_vector<clanguml::common::model::package> packages_;
 };
diff --git a/src/sequence_diagram/model/diagram.cc b/src/sequence_diagram/model/diagram.cc
index 89c75692..4e73c9ba 100644
--- a/src/sequence_diagram/model/diagram.cc
+++ b/src/sequence_diagram/model/diagram.cc
@@ -45,6 +45,18 @@ std::string diagram::to_alias(const std::string &full_name) const
     return full_name;
 }
 
+inja::json diagram::context() const {
+    inja::json ctx;
+    ctx["name"] = name();
+    ctx["type"] = "sequence";
+
+    inja::json::array_t elements{};
+
+    ctx["elements"] = elements;
+
+    return ctx;
+}
+
 }
 
 namespace clanguml::common::model {
diff --git a/src/sequence_diagram/model/diagram.h b/src/sequence_diagram/model/diagram.h
index 808e0899..9190b12e 100644
--- a/src/sequence_diagram/model/diagram.h
+++ b/src/sequence_diagram/model/diagram.h
@@ -45,6 +45,8 @@ public:
 
     std::string to_alias(const std::string &full_name) const;
 
+    inja::json context() const override;
+
     bool started{false};
 
     std::map<int64_t, activity> sequences;
diff --git a/tests/test_cases.cc b/tests/test_cases.cc
index 5cf6c2b6..e06c72a4 100644
--- a/tests/test_cases.cc
+++ b/tests/test_cases.cc
@@ -242,16 +242,17 @@ using namespace clanguml::test::matchers;
 #include "t00047/test_case.h"
 #include "t00048/test_case.h"
 #include "t00049/test_case.h"
+#include "t00050/test_case.h"
 
-////
-//// Sequence diagram tests
-////
+///
+/// Sequence diagram tests
+///
 #include "t20001/test_case.h"
 #include "t20002/test_case.h"
 
-////
-//// Package diagram tests
-////
+///
+/// Package diagram tests
+///
 #include "t30001/test_case.h"
 #include "t30002/test_case.h"
 #include "t30003/test_case.h"
@@ -261,21 +262,21 @@ using namespace clanguml::test::matchers;
 #include "t30007/test_case.h"
 #include "t30008/test_case.h"
 
-////
-//// Include diagram tests
-////
+///
+/// Include diagram tests
+///
 #include "t40001/test_case.h"
 #include "t40002/test_case.h"
 #include "t40003/test_case.h"
 
-////
-//// Other tests (e.g. configuration file)
-////
+///
+/// Other tests (e.g. configuration file)
+///
 #include "t90000/test_case.h"
 
-//
-// Main test function
-//
+///
+/// Main test function
+///
 int main(int argc, char *argv[])
 {
     Catch::Session session;
diff --git a/tests/test_cases.h b/tests/test_cases.h
index 84aae50f..9fe8e489 100644
--- a/tests/test_cases.h
+++ b/tests/test_cases.h
@@ -367,7 +367,7 @@ ContainsMatcher IsLayoutHint(std::string const &from, std::string const &hint,
 }
 
 ContainsMatcher HasNote(std::string const &cls, std::string const &position,
-    std::string const &note,
+    std::string const &note = "",
     CaseSensitive::Choice caseSensitivity = CaseSensitive::Yes)
 {
     return ContainsMatcher(CasedString(
diff --git a/tests/test_cases.yaml b/tests/test_cases.yaml
index 1ca8818f..12be8585 100644
--- a/tests/test_cases.yaml
+++ b/tests/test_cases.yaml
@@ -144,6 +144,9 @@ test_cases:
     - name: t00049
       title: Test case configurable type aliases
       description:
+    - name: t00050
+      title: Test case for generating notes from comments using jinja templates
+      description:
   Sequence diagrams:
     - name: t20001
       title: Basic sequence diagram test case