diff --git a/lib/providers/collection_providers.dart b/lib/providers/collection_providers.dart
index 5c867675..c6e9db55 100644
--- a/lib/providers/collection_providers.dart
+++ b/lib/providers/collection_providers.dart
@@ -3,7 +3,6 @@ import 'package:apidash_core/apidash_core.dart';
 import 'package:flutter/material.dart';
 import 'package:flutter_riverpod/flutter_riverpod.dart';
 import 'package:apidash/consts.dart';
-import 'package:genai/genai.dart';
 import 'providers.dart';
 import '../models/models.dart';
 import '../services/services.dart';
diff --git a/lib/widgets/ai_ui_desginer_widgets.dart b/lib/screens/home_page/editor_pane/agentic_ui_elements/ai_ui_designer/ai_ui_desginer_widgets.dart
similarity index 94%
rename from lib/widgets/ai_ui_desginer_widgets.dart
rename to lib/screens/home_page/editor_pane/agentic_ui_elements/ai_ui_designer/ai_ui_desginer_widgets.dart
index 06a0ff55..5c77758c 100644
--- a/lib/widgets/ai_ui_desginer_widgets.dart
+++ b/lib/screens/home_page/editor_pane/agentic_ui_elements/ai_ui_designer/ai_ui_desginer_widgets.dart
@@ -1,5 +1,6 @@
 import 'dart:convert';
 import 'package:apidash/consts.dart';
+import 'package:apidash/providers/collection_providers.dart';
 import 'package:apidash/services/agentic_services/agent_caller.dart';
 import 'package:apidash/widgets/widget_sending.dart';
 import 'package:apidash_design_system/apidash_design_system.dart';
@@ -8,7 +9,7 @@ import 'package:flutter/services.dart';
 import 'package:flutter_riverpod/flutter_riverpod.dart';
 import 'package:genai/agentic_engine/blueprint.dart';
 import 'package:stac/stac.dart' as stac;
-import '../services/agentic_services/agents/agents.dart';
+import '../../../../../services/agentic_services/agents/agents.dart';
 
 void showCustomDialog(BuildContext context, Widget dialogContent) {
   showDialog(
@@ -497,3 +498,33 @@ class StacRenderer extends StatelessWidget {
     );
   }
 }
+
+class AIGenerateUIButton extends ConsumerWidget {
+  const AIGenerateUIButton({super.key});
+
+  @override
+  Widget build(BuildContext context, WidgetRef ref) {
+    return FilledButton.tonalIcon(
+      style: FilledButton.styleFrom(
+        padding: kPh12,
+        minimumSize: const Size(44, 44),
+      ),
+      onPressed: () {
+        final model = ref.watch(selectedRequestModelProvider
+            .select((value) => value?.httpResponseModel));
+        showCustomDialog(
+          context,
+          GenerateUIDialog(content: model?.formattedBody ?? ""),
+        );
+      },
+      icon: Icon(
+        Icons.generating_tokens,
+      ),
+      label: const SizedBox(
+        child: Text(
+          kLabelGenerateUI,
+        ),
+      ),
+    );
+  }
+}
diff --git a/lib/widgets/ai_toolgen_widgets.dart b/lib/screens/home_page/editor_pane/agentic_ui_elements/tool_generation/ai_toolgen_widgets.dart
similarity index 94%
rename from lib/widgets/ai_toolgen_widgets.dart
rename to lib/screens/home_page/editor_pane/agentic_ui_elements/tool_generation/ai_toolgen_widgets.dart
index 9da8df7d..9c064c58 100644
--- a/lib/widgets/ai_toolgen_widgets.dart
+++ b/lib/screens/home_page/editor_pane/agentic_ui_elements/tool_generation/ai_toolgen_widgets.dart
@@ -7,7 +7,7 @@ import 'package:apidash/screens/common_widgets/ai/ai_model_selector_button.dart'
 import 'package:apidash/services/agentic_services/agent_caller.dart';
 import 'package:apidash/services/agentic_services/agents/apitool_bodygen.dart';
 import 'package:apidash/services/agentic_services/agents/apitool_funcgen.dart';
-import 'package:apidash/widgets/ai_ui_desginer_widgets.dart';
+import 'package:apidash/screens/home_page/editor_pane/agentic_ui_elements/ai_ui_designer/ai_ui_desginer_widgets.dart';
 import 'package:apidash/widgets/button_copy.dart';
 import 'package:apidash/widgets/previewer_code.dart';
 import 'package:apidash/widgets/widget_sending.dart';
@@ -18,7 +18,7 @@ import 'package:flutter/material.dart';
 import 'package:flutter_riverpod/flutter_riverpod.dart';
 import 'package:genai/agentic_engine/blueprint.dart';
 
-import '../providers/providers.dart';
+import '../../../../../providers/providers.dart';
 
 class GenerateToolDialog extends ConsumerStatefulWidget {
   final APIDashRequestDescription requestDesc;
@@ -426,3 +426,28 @@ class DefaultLLModelSelectorWidget extends ConsumerWidget {
     );
   }
 }
+
+class GenerateToolButton extends ConsumerWidget {
+  const GenerateToolButton({super.key});
+
+  @override
+  Widget build(BuildContext context, WidgetRef ref) {
+    return FilledButton.tonalIcon(
+      style: FilledButton.styleFrom(
+        padding: kPh12,
+        minimumSize: const Size(44, 44),
+      ),
+      onPressed: () async {
+        GenerateToolDialog.show(context, ref);
+      },
+      icon: Icon(
+        Icons.token_outlined,
+      ),
+      label: const SizedBox(
+        child: Text(
+          "Generate Tool",
+        ),
+      ),
+    );
+  }
+}
diff --git a/lib/widgets/response_body_success.dart b/lib/widgets/response_body_success.dart
index 67d47777..b31a0971 100644
--- a/lib/widgets/response_body_success.dart
+++ b/lib/widgets/response_body_success.dart
@@ -1,10 +1,5 @@
-import 'dart:convert';
-
-import 'package:apidash/apitoolgen/request_consolidator.dart';
-import 'package:apidash/providers/collection_providers.dart';
-import 'package:apidash/widgets/ai_toolgen_widgets.dart';
-import 'package:apidash/widgets/ai_ui_desginer_widgets.dart';
-
+import 'package:apidash/screens/home_page/editor_pane/agentic_ui_elements/tool_generation/ai_toolgen_widgets.dart';
+import 'package:apidash/screens/home_page/editor_pane/agentic_ui_elements/ai_ui_designer/ai_ui_desginer_widgets.dart';
 import 'package:apidash_core/apidash_core.dart';
 import 'package:apidash_design_system/apidash_design_system.dart';
 import 'package:flutter/foundation.dart';
@@ -12,11 +7,9 @@ import 'package:flutter/material.dart';
 import 'package:apidash/utils/utils.dart';
 import 'package:apidash/widgets/widgets.dart';
 import 'package:apidash/consts.dart';
-import 'package:genai/genai.dart';
-import 'package:flutter_riverpod/flutter_riverpod.dart';
 import 'button_share.dart';
 
-class ResponseBodySuccess extends ConsumerStatefulWidget {
+class ResponseBodySuccess extends StatefulWidget {
   const ResponseBodySuccess({
     super.key,
     required this.mediaType,
@@ -42,11 +35,10 @@ class ResponseBodySuccess extends ConsumerStatefulWidget {
   final bool isPartOfHistory;
 
   @override
-  ConsumerState<ResponseBodySuccess> createState() =>
-      _ResponseBodySuccessState();
+  State<ResponseBodySuccess> createState() => _ResponseBodySuccessState();
 }
 
-class _ResponseBodySuccessState extends ConsumerState<ResponseBodySuccess> {
+class _ResponseBodySuccessState extends State<ResponseBodySuccess> {
   int segmentIdx = 0;
 
   @override
@@ -77,46 +69,9 @@ class _ResponseBodySuccessState extends ConsumerState<ResponseBodySuccess> {
                 Row(
                   mainAxisAlignment: MainAxisAlignment.end,
                   children: [
-                    FilledButton.tonalIcon(
-                      style: FilledButton.styleFrom(
-                        padding: kPh12,
-                        minimumSize: const Size(44, 44),
-                      ),
-                      onPressed: () async {
-                        GenerateToolDialog.show(context, ref);
-                      },
-                      icon: Icon(
-                        Icons.token_outlined,
-                      ),
-                      label: const SizedBox(
-                        child: Text(
-                          "Generate Tool",
-                        ),
-                      ),
-                    ),
+                    GenerateToolButton(),
                     kHSpacer10,
-                    FilledButton.tonalIcon(
-                      style: FilledButton.styleFrom(
-                        padding: kPh12,
-                        minimumSize: const Size(44, 44),
-                      ),
-                      onPressed: () {
-                        final model = ref.watch(selectedRequestModelProvider
-                            .select((value) => value?.httpResponseModel));
-                        showCustomDialog(
-                          context,
-                          GenerateUIDialog(content: model?.formattedBody ?? ""),
-                        );
-                      },
-                      icon: Icon(
-                        Icons.generating_tokens,
-                      ),
-                      label: const SizedBox(
-                        child: Text(
-                          kLabelGenerateUI,
-                        ),
-                      ),
-                    ),
+                    AIGenerateUIButton(),
                     kHSpacer10,
                   ],
                 ),