diff --git a/Classes/ViewHelpers/SkillGroupingViewHelper.php b/Classes/ViewHelpers/SkillGroupingViewHelper.php
index 884903e..1701ace 100644
--- a/Classes/ViewHelpers/SkillGroupingViewHelper.php
+++ b/Classes/ViewHelpers/SkillGroupingViewHelper.php
@@ -22,8 +22,8 @@ namespace SkillDisplay\SkilldisplayContent\ViewHelpers;
  */
 
 use Closure;
-use Exception;
-use SkillDisplay\PHPToolKit\Verification\Link;
+use SkillDisplay\PHPToolKit\Entity\Skill;
+use SkillDisplay\PHPToolKit\Entity\SkillSet;
 use TYPO3Fluid\Fluid\Core\Rendering\RenderingContextInterface;
 use TYPO3Fluid\Fluid\Core\ViewHelper\AbstractViewHelper;
 use TYPO3Fluid\Fluid\Core\ViewHelper\Traits\CompileWithRenderStatic;
@@ -47,7 +47,7 @@ class SkillGroupingViewHelper extends AbstractViewHelper
     public function initializeArguments()
     {
         $this->registerArgument('skills', 'array', 'An array of skills to group', false, []);
-        $this->registerArgument('skillSet', 'array', 'The skills of this skill set will be grouped', false, []);
+        $this->registerArgument('skillSet', SkillSet::class, 'The skills of this skill set will be grouped', false, []);
         $this->registerArgument('as', 'string', 'The name of the iteration variable', true);
     }
 
@@ -56,19 +56,21 @@ class SkillGroupingViewHelper extends AbstractViewHelper
         Closure $renderChildrenClosure,
         RenderingContextInterface $renderingContext
     ) {
+        /** @var Skill[] $skills */
         $skills = $arguments['skills'];
+        /** @var SkillSet $skillSet */
         $skillSet = $arguments['skillSet'];
         $as = $arguments['as'];
 
         if ($skillSet) {
-            $groups = self::group($skillSet['skills']);
+            $groups = self::group($skillSet->getSkills());
         } else {
             $groups = self::group($skills);
         }
 
         $templateVariableContainer = $renderingContext->getVariableProvider();
         $output = '';
-        foreach ($groups['values'] as $groups => $group) {
+        foreach ($groups as $group) {
             $templateVariableContainer->add($as, $group);
             $output .= $renderChildrenClosure();
             $templateVariableContainer->remove($as);
@@ -76,6 +78,10 @@ class SkillGroupingViewHelper extends AbstractViewHelper
         return $output;
     }
 
+    /**
+     * @param Skill[] $skills
+     * @return array
+     */
     protected static function group(array $skills): array
     {
         $result = [