Update app.py

app.py

@@ -3,13 +3,13 @@ import google.generativeai as genai
 import re
 import datetime
 import os
-import ipaddress
+import yaml
 
 # Secret key and Google Gemini API configuration
 API_KEY = st.secrets["GOOGLE_API_KEY"]
 
 # Page configuration
-st.set_page_config(page_title="⚙️ Gemini2 Ansible Gen Pro", page_icon="⚙️", layout="wide")
+st.set_page_config(page_title="⚙️ Gemini2 Puppet Bolt Gen Pro", page_icon="⚙️", layout="wide")
 
 # --- Helper Functions ---
 def send_message_to_model(message, model_name, temperature, top_p, top_k, max_tokens):
@@ -34,25 +34,27 @@ def send_message_to_model(message, model_name, temperature, top_p, top_k, max_to
         st.error(f"❌ Error communicating with the AI: {e}")
         return None
 
-def generate_ansible_playbook(prompt_base, detail_level, os_type, security_level, model_name, temperature, top_p, top_k, max_tokens, prompt_detail, encoding, add_header, log_level, ansible_user, ansible_ssh_pass, ansible_become, custom_requirements, specific_details, target_hosts_description):
-    """Generates an Ansible playbook based on user settings."""
+def generate_puppet_bolt_plan(prompt_base, detail_level, os_type, security_level, model_name, temperature, top_p, top_k, max_tokens, prompt_detail, encoding, add_header, log_level, bolt_user, bolt_ssh_pass, custom_requirements, specific_details, target_nodes_description, task_type, task_code, parameters_description):
+    """Generates a Puppet Bolt plan based on user settings."""
     prompt = f"""
-    You are an expert Ansible automation engineer. Your task is to generate a complete, secure, and efficient Ansible playbook based on the following description:
+    You are an expert automation engineer specializing in Puppet Bolt. Your task is to generate a complete, secure, and efficient Puppet Bolt plan based on the following description:
 
-    **Goal:** Create the most complete, detailed, efficient, and secure Ansible playbook possible, considering all variables, edge cases, and potential scenarios.
+    **Goal:** Create the most complete, detailed, efficient, and secure Puppet Bolt plan possible, considering all variables, edge cases, and potential scenarios.
 
-    **Playbook Description:** {prompt_base}
+    **Plan Description:** {prompt_base}
 
     **Detail Level:** {detail_level}
     **Security Level:** {security_level}
     **Prompt Detail Level**:{prompt_detail}
     **Operating System:** {os_type}
-    **Target Hosts:** {target_hosts_description if target_hosts_description else "None"}
-    **Ansible User:** {ansible_user if ansible_user else "root"}
-    **Ansible SSH Password:** {ansible_ssh_pass if ansible_ssh_pass else "None"}
-     **Ansible Become:** {ansible_become}
+    **Target Nodes:** {target_nodes_description if target_nodes_description else "None"}
+    **Bolt User:** {bolt_user if bolt_user else "root"}
+    **Bolt SSH Password:** {bolt_ssh_pass if bolt_ssh_pass else "None"}
+     **Parameters:** {parameters_description if parameters_description else "None"}
     **Custom Requirements:** {custom_requirements if custom_requirements else "None"}
     **Specific Details:** {specific_details if specific_details else "None"}
+    **Task Type:** {task_type}
+    **Task Code:** {task_code if task_code else "None"}
 
     **Response Format:**
     - Respond in Markdown format, including a YAML code block with its original formatting, without line breaks.
@@ -60,45 +62,50 @@ def generate_ansible_playbook(prompt_base, detail_level, os_type, security_level
     - Do not include comments, explanations, or any other text outside the code block, unless the prompt requires.
     - The code must maintain its full vertical formatting, respecting indentation and line breaks.
     - The code must be realistic, using real-world examples, data, and situations.
-     - The `hosts` section of the playbook, must include the list of the target hosts in the description, if the target hosts are not specified you must create your own hosts list, please avoid using "all" if the target hosts are specified.
+     - The plan must use advanced Bolt resources, such as tasks, parameters, variables, functions, conditionals, loops and targets when necessary.
+      -The plan must use the `run_task` function to execute the tasks.
+      -The plan must use a valid node target.
     - Explore different approaches, techniques, and advanced practices, always prioritizing security and efficiency.
-    - Use advanced Ansible resources such as modules, roles, variables, functions, conditionals, loops, and handlers, when necessary.
-    - Unless the user specifies otherwise, use the most current and secure versions of the modules, using ansible and following best practices.
-    - If the description asks to create a file, the ansible playbook should create the file directly in the file system and not use a screen output for this.
-    - If the description asks to read a file, the ansible playbook should read the file directly from the file system and not use a screen input for this.
-     - If possible, use Jinja2 templates when required and do not create hard code in your code
+     - The plan must use native Bolt functions and modules instead of others.
+    - Unless the user specifies otherwise, use the most current and secure versions of the modules, using bolt and following best practices.
+    - If the description asks to create a file, the plan should create the file directly in the file system and not use a screen output for this.
+    - If the description asks to read a file, the plan should read the file directly from the file system and not use a screen input for this.
     - Use incremental reasoning to add improvements, expansions, and considerations to your code.
     - Use the history of the conversations so that the response is incremental.
 
     **Log Level**:{log_level}
 
     **Important:**
-    - Generate only one playbook at a time.
+    - Generate only one plan at a time.
     - Create the longest, most complete, and detailed code possible to cover a wide range of possibilities and scenarios.
     - Consider all the details of the request, expanding the response and improving the configuration.
-     - Use the `delegate_to`, `local_action` and `register` options when necessary.
-     - Use the `with_items` and `with_dict` options when necessary.
-     - Do not use `sudo` unless absolutely required, prefer `become`
-     -  If possible, use `ansible.builtin` modules instead of other modules.
-
+     - If the prompt does not include parameters, but if the user add a variable in a code block, create a task parameter to use that variable.
+     - If the prompt includes parameters, create a task parameter to use that parameters.
     """
-    response = send_message_to_model(prompt, model_name, temperature, top_p, top_k, max_tokens)
+    response = send_message_to_model(message, model_name, temperature, top_p, top_k, max_tokens)
     return response
 
 def parse_and_save_yml(ai_code, short_title, encoding, add_header):
-    """Parses the markdown and saves the ansible playbook as .yml."""
+    """Parses the markdown and saves the Puppet Bolt plan as .yml."""
     match = re.search(r'```yaml\s*(.*?)\s*```', ai_code, re.DOTALL | re.IGNORECASE)
     if match:
         yml_code = match.group(1).strip()
     else:
         yml_code = ai_code.strip()
-
-    file_name = f"ansible_{short_title}.yml"
+        
+    try:
+        yaml.safe_load(yml_code)
+    except yaml.YAMLError as e:
+        st.error(f"❌ Generated YAML is invalid: {e}")
+        return None, None
+    
+    
+    file_name = f"bolt_{short_title}.yml"
 
     if add_header:
         current_date = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         header = f"""#===============================================================================
-# Ansible Playbook Generated by Google Gemini 2 Ansible Gen Pro
+# Puppet Bolt Plan Generated by Google Gemini 2 Bolt Gen Pro
 # Date: {current_date}
 # Author: Elias Andrade AKA Chaos4455
 #===============================================================================\n"""
@@ -111,8 +118,8 @@ def parse_and_save_yml(ai_code, short_title, encoding, add_header):
 
 
 def main():
-    st.title("⚙️ Gemini2 Ansible Gen Pro by [Elias Andrade](https://github.com/chaos4455)")
-    st.markdown("Generate advanced Ansible playbooks with ease! 🚀")
+    st.title("⚙️ Gemini2 Puppet Bolt Gen Pro by [Elias Andrade](https://github.com/chaos4455)")
+    st.markdown("Generate advanced Puppet Bolt plans with ease! 🚀")
     st.markdown("---")
 
     # Layout in columns (sidebar and main area)
@@ -129,49 +136,56 @@ def main():
             max_tokens = st.number_input("📏 Max Tokens", min_value=128, max_value=8192, value=8192, step=128, help="Adjust the maximum size of the response.")
             
         with st.expander("📝 Prompt Settings"):
-            prompt_presets = st.selectbox("🎯 Predefined Prompts", ["None", "Install packages", "Configure services", "Manage users", "File management", "System updates", "Deploy applications"], index=0, help="Choose a predefined prompt.")
-            prompt_detail = st.selectbox("🧐 Prompt Detail", ["More descriptive", "Default", "Concise"], index=1, help="Defines the level of detail of the prompt")
+             prompt_presets = st.selectbox("🎯 Predefined Prompts", ["None", "Install packages", "Run commands", "Manage services", "Manage files", "System Updates", "Copy Files"], index=0, help="Choose a predefined prompt.")
+             prompt_detail = st.selectbox("🧐 Prompt Detail", ["More descriptive", "Default", "Concise"], index=1, help="Defines the level of detail of the prompt")
 
-        with st.expander("⚙️ Ansible Settings"):
+        with st.expander("⚙️ Puppet Bolt Settings"):
             os_type = st.selectbox(
                 "💻 Operating System",
                 [
-                    "Ubuntu", "CentOS", "RedHat", "Debian", "Windows", "Other"
+                    "Linux", "Windows", "Other"
                 ],
                 index=0, help="Choose the target operating system."
             )
             
-            ansible_user = st.text_input("👤 Ansible User", value="root", help="User for Ansible SSH connection")
-            ansible_ssh_pass = st.text_input("🔑 Ansible SSH Password", type="password", help="Password for Ansible SSH connection (Use with caution)")
-            ansible_become = st.checkbox("🔥 Use Become", value=True, help="Enable privilege escalation")
+            bolt_user = st.text_input("👤 Bolt User", value="root", help="User for Bolt SSH connection")
+            bolt_ssh_pass = st.text_input("🔑 Bolt SSH Password", type="password", help="Password for Bolt SSH connection (Use with caution)")
 
             encoding = st.selectbox("🔤 Encoding", ["utf-8", "ansi"], index=0, help="Choose the encoding of the .yml file.")
             add_header = st.checkbox("📜 Add Header", value=True, help="Add a header with information in the .yml file.")
             log_level = st.selectbox("🗂️ Logging Level", ["Detailed", "Default", "Minimum"], index=1, help="Defines the detail level of logs.")
             security_level = st.radio("Security Level", ["High", "Medium", "Low"], index=1)
-            custom_requirements = st.text_input("Custom Requirements", placeholder="Ex: Use specific modules, parameters", help="Add specific requirements for the generated playbook.")
-            specific_details = st.text_input("Specific Details", placeholder="Ex: Specific configurations, edge cases", help="Add specific details for playbook generation.")
+            custom_requirements = st.text_input("Custom Requirements", placeholder="Ex: Use specific modules, parameters", help="Add specific requirements for the generated plan.")
+            specific_details = st.text_input("Specific Details", placeholder="Ex: Specific configurations, edge cases", help="Add specific details for plan generation.")
             
             detail_level = st.selectbox("Detail Level", ["More detailed", "Default", "More concise"], index=1)
-
+           
+        with st.expander("🔨 Task Settings"):
+            task_type = st.selectbox("🎯 Task Type", ["command", "script", "puppet"], index = 0, help="Type of task")
+            task_code = st.text_area("📋 Task Code",  placeholder="Command or Code for the task", help="Enter your command, script or Puppet code")
+            parameters_description = st.text_input("🧰 Parameters Description",  placeholder="List of parameters for the task", help="Enter the parameters needed for this task")
 
     with col2:
         # User's base prompt
-        prompt_base = st.text_input("Describe the Ansible Playbook:", placeholder="Ex: Install apache2 and configure a default page for all the servers", key="prompt_base")
-        target_hosts_description = st.text_input("Target Hosts Description (text format, IP(s), Range or CIDR):", placeholder="Ex: all the servers in 192.168.0.0/24 network",key="target_hosts")
+        prompt_base = st.text_input("Describe the Puppet Bolt Plan:", placeholder="Ex: Execute a command to check disk space on all Linux servers, and save the output on /tmp/output.txt", key="prompt_base")
+        target_nodes_description = st.text_input("Target Nodes Description (text format, IP(s), Range or CIDR):", placeholder="Ex: all servers in the 192.168.0.0/24 network",key="target_nodes")
 
         if prompt_presets != "None":
             if prompt_base:
                prompt_base = f"{prompt_presets} , {prompt_base}"
             else:
                prompt_base = prompt_presets;
-        if st.button("✨ Generate Ansible Playbook"):
+        if st.button("✨ Generate Puppet Bolt Plan"):
             if not prompt_base:
-                st.error("⚠️ Please enter a playbook description.")
+                st.error("⚠️ Please enter a plan description.")
                 return
-
-            with st.spinner("⏳ Generating playbook..."):
-                ai_code = generate_ansible_playbook(
+            
+            if task_type != "command" and not task_code:
+                 st.error("⚠️ Please enter the code for the task.")
+                 return
+            
+            with st.spinner("⏳ Generating plan..."):
+                ai_code = generate_puppet_bolt_plan(
                    prompt_base, 
                    detail_level, 
                    os_type,
@@ -185,29 +199,32 @@ def main():
                    encoding,
                    add_header,
                    log_level,
-                   ansible_user,
-                   ansible_ssh_pass,
-                   ansible_become,
+                   bolt_user,
+                   bolt_ssh_pass,
                    custom_requirements,
                    specific_details,
-                   target_hosts_description # Pass only the description
+                   target_nodes_description,
+                   task_type,
+                   task_code,
+                   parameters_description
                    )
 
                 if ai_code:
-                    st.markdown("### ✅ Generated Playbook:")
+                    st.markdown("### ✅ Generated Plan:")
                     st.code(ai_code, language="yaml")
 
                     short_title = prompt_base[:30].strip().replace(" ", "_").lower()
                     file_name_yml, yml_code = parse_and_save_yml(ai_code, short_title, encoding, add_header)
-
-                    st.download_button(
-                        label="⬇️ Download Playbook (.yml)",
-                        data=yml_code,
-                        file_name=file_name_yml,
-                        mime="application/x-yaml",
-                    )
+                    
+                    if file_name_yml and yml_code: #check if the file was generated
+                        st.download_button(
+                            label="⬇️ Download Plan (.yml)",
+                            data=yml_code,
+                            file_name=file_name_yml,
+                            mime="application/x-yaml",
+                        )
                 else:
-                    st.error("❌ Error generating the playbook. Check the connection with the AI and try again.")
+                    st.error("❌ Error generating the plan. Check the connection with the AI and try again.")
 
 if __name__ == "__main__":
     main()