diff --git a/src/Generators/stubs/builder.stub b/src/Generators/stubs/builder.stub
index 2ed2af7..511be79 100644
--- a/src/Generators/stubs/builder.stub
+++ b/src/Generators/stubs/builder.stub
@@ -13,7 +13,7 @@ class DummyClass extends DataTable
     /**
      * Build the DataTable class.
      *
-     * @param QueryBuilder $query Results from query() method.
+     * @param QueryBuilder<ModelName> $query Results from query() method.
      */
     public function dataTable(QueryBuilder $query): EloquentDataTable
     {
@@ -24,6 +24,8 @@ class DummyClass extends DataTable
 
     /**
      * Get the query source of dataTable.
+     *
+     * @return QueryBuilder<ModelName>
      */
     public function query(ModelName $model): QueryBuilder
     {
diff --git a/src/Generators/stubs/datatables.stub b/src/Generators/stubs/datatables.stub
index bccfede..1a574c1 100644
--- a/src/Generators/stubs/datatables.stub
+++ b/src/Generators/stubs/datatables.stub
@@ -17,7 +17,7 @@ class DummyClass extends DataTable
     /**
      * Build the DataTable class.
      *
-     * @param QueryBuilder $query Results from query() method.
+     * @param QueryBuilder<ModelName> $query Results from query() method.
      */
     public function dataTable(QueryBuilder $query): EloquentDataTable
     {
@@ -28,6 +28,8 @@ class DummyClass extends DataTable
 
     /**
      * Get the query source of dataTable.
+     *
+     * @return QueryBuilder<ModelName>
      */
     public function query(ModelName $model): QueryBuilder
     {