diff --git a/src/Analysis.php b/src/Analysis.php
index 6ede8d90..a200e402 100644
--- a/src/Analysis.php
+++ b/src/Analysis.php
@@ -323,7 +323,19 @@ public function getAnnotationsOfType($classes, bool $strict = false): array
     /**
      * @param string $fqdn the source class/interface/trait
      */
-    public function getSchemaForSource(string $fqdn, string $class = OA\Schema::class): ?OA\AbstractAnnotation
+    public function getSchemaForSource(string $fqdn): ?OA\Schema
+    {
+        return $this->getAnnotationForSource($fqdn, OA\Schema::class);
+    }
+
+    /**
+     * @template T of OA\AbstractAnnotation
+     *
+     * @param  string          $fqdn  the source class/interface/trait
+     * @param  class-string<T> $class
+     * @return T|null
+     */
+    public function getAnnotationForSource(string $fqdn, string $class): ?OA\AbstractAnnotation
     {
         $fqdn = '\\' . ltrim($fqdn, '\\');
 
@@ -331,6 +343,7 @@ public function getSchemaForSource(string $fqdn, string $class = OA\Schema::clas
             if (array_key_exists($fqdn, $definitions)) {
                 $definition = $definitions[$fqdn];
                 if (is_iterable($definition['context']->annotations)) {
+                    /** @var OA\AbstractAnnotation $annotation */
                     foreach (array_reverse($definition['context']->annotations) as $annotation) {
                         if (is_a($annotation, $class) && $annotation->isRoot($class) && !$annotation->_context->is('generated')) {
                             return $annotation;
diff --git a/src/Processors/AugmentRefs.php b/src/Processors/AugmentRefs.php
index 1c603676..b827efa1 100644
--- a/src/Processors/AugmentRefs.php
+++ b/src/Processors/AugmentRefs.php
@@ -66,7 +66,7 @@ protected function resolveFQCNRefs(Analysis $analysis): void
                 // check if we have a schema for this
                 if ($refSchema = $analysis->getSchemaForSource($annotation->ref)) {
                     $annotation->ref = OA\Components::ref($refSchema);
-                } elseif ($refSchema = $analysis->getSchemaForSource($annotation->ref, get_class($annotation))) {
+                } elseif ($refSchema = $analysis->getAnnotationForSource($annotation->ref, get_class($annotation))) {
                     $annotation->ref = OA\Components::ref($refSchema);
                 }
             }