diff --git a/util/prowler_check_kreator/lib/llms/gemini.py b/util/prowler_check_kreator/lib/llms/gemini.py
index 51334da6337..830e56bc6d0 100644
--- a/util/prowler_check_kreator/lib/llms/gemini.py
+++ b/util/prowler_check_kreator/lib/llms/gemini.py
@@ -1,5 +1,6 @@
 import json
 import os
+import re
 
 import google.generativeai as genai
 
@@ -67,7 +68,11 @@ def _generate_content(self, prompt_parts: list) -> str:
 
     def _prepare_check_prompt(self, check_name: str, check_reference: str) -> list:
         """Prepare the prompt for generating the check."""
-        class_name = check_reference.split("(")[0].split("class ")[1]
+
+        class_name = re.search(
+            r"class\s+([A-Za-z_][A-Za-z0-9_]*)\s*\((.*?)\)\s*:", check_reference
+        ).group(1)
+
         prompt_parts = [
             f"Your task is to create a new security check called '{check_name}' for Prowler (a Cloud Security tool).",
             "The control is a Python class that inherits from the Check class and has only one method called execute.",
diff --git a/util/prowler_check_kreator/prowler_check_kreator.py b/util/prowler_check_kreator/prowler_check_kreator.py
index 4795d12022e..3ed08e485f2 100644
--- a/util/prowler_check_kreator/prowler_check_kreator.py
+++ b/util/prowler_check_kreator/prowler_check_kreator.py
@@ -83,7 +83,7 @@ def __init__(self, provider: str, check_name: str):
         # Let the user to use the model that he wants
         self._model = None
         supported_models = [
-            "gemini-1.5-flash (default)",
+            "gemini-1.5-flash",
             "gemini-1.5-pro",
             "gemini-1.0-pro",
         ]
@@ -361,8 +361,8 @@ def _write_test_file(self) -> None:
         sys.exit(prowler_check_creator.kreate_check())
 
     except ValueError as e:
-        print(f"Error: {e}")
+        print(f"Error: {e} (line {sys.exc_info()[-1].tb_lineno})")
         sys.exit(1)
     except Exception as e:
-        print(f"Error: {e}")
+        print(f"Error: {e} (line {sys.exc_info()[-1].tb_lineno})")
         sys.exit(1)