main.py

from langchain_google_genai import ChatGoogleGenerativeAI
# from langchain_groq import ChatGroq
from windows_use.agent import Agent
from dotenv import load_dotenv

from prompt_generator import generate_user_prompt
import time
import subprocess
import psutil

load_dotenv()

def get_process_to_kill(prompt: str) -> str | None:
    """Maps a prompt keyword to a process executable name."""
    prompt = prompt.lower()

    process_map = {
        "zoom": "Zoom.exe",
        "chrome": "chrome.exe",
        "microsoft edge": "msedge.exe",
        "firefox": "firefox.exe",
        "microsoft teams": "ms-teams.exe", 
        "discord": "Discord.exe",
        "slack": "slack.exe",
        "notepad": "notepad.exe"
    }

    for keyword, process_name in process_map.items():
        if keyword in prompt:
            return process_name
    return None

def kill_process(process_name: str):
    """Forcefully terminates a process by its name using taskkill."""
    if process_name:
        print(f"Attempting to close process: {process_name}...")
        try:
            command = ["taskkill", "/F", "/IM", process_name, "/T"]
            subprocess.run(command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            print(f"Successfully closed {process_name}.")
        except (subprocess.CalledProcessError, FileNotFoundError):
            print(f"Could not close {process_name}. It may not have been running.")
    else:
        print("No specific process identified to close for this task.")

def main():
    llm=ChatGoogleGenerativeAI(model='gemini-2.0-flash')
    # llm=ChatGroq(model='meta-llama/llama-4-scout-17b-16e-instruct',api_key=os.getenv("GROQ_API_KEY"))
    agent = Agent(llm=llm,browser='chrome',use_vision=False)

    # query=input("Enter your query: ")
    # agent.print_response(query)

    for i in range(2):
        query = generate_user_prompt()
        print(f"\n[Run {i+1}] Prompt: {query}")

        agent.print_response(query)

        print("Waiting for application to stabilize before closing...")
        time.sleep(5) # 5-second delay

        process_to_kill = get_process_to_kill(query)
        kill_process(process_to_kill)

        time.sleep(5) # 5 second delay before the next prompt

if __name__ == "__main__":
    main()