Update command_classifier.py

command_classifier.py

@@ -1,11 +1,46 @@
 def classify_command(command: str) -> str:
+    """
+    Classifies a user's command into categories for action or AI response.
+
+    Args:
+        command (str): User's input text.
+
+    Returns:
+        str: Command type (e.g. 'open_app', 'screenshot', 'create_file', etc.)
+    """
     command = command.lower()
-    if "open" in command and any(app in command for app in ["chrome", "firefox", "notepad", "calculator"]):
+
+    # --- Application control ---
+    if "open" in command and any(app in command for app in ["chrome", "firefox", "notepad", "calculator", "edge", "terminal", "vscode"]):
         return "open_app"
-    elif "screenshot" in command:
-        return "screenshot"
-    elif "create" in command and "file" in command:
+    elif "close" in command and any(app in command for app in ["chrome", "firefox", "notepad", "calculator", "edge", "terminal", "vscode"]):
+        return "close_app"
+
+    # --- File system operations ---
+    if "create" in command and "file" in command:
         return "create_file"
-    elif "search" in command or "what is" in command:
+    elif "delete" in command and "file" in command:
+        return "delete_file"
+    elif "read" in command and "file" in command:
+        return "read_file"
+    elif "write" in command and "file" in command:
+        return "write_file"
+
+    # --- System commands ---
+    if "screenshot" in command or "screen shot" in command:
+        return "screenshot"
+    if "shutdown" in command:
+        return "shutdown"
+    if "restart" in command:
+        return "restart"
+    if "lock" in command:
+        return "lock_computer"
+
+    # --- Internet-related ---
+    if "search" in command or "what is" in command or "who is" in command:
         return "web_search"
+    if "open" in command and ".com" in command:
+        return "open_website"
+
+    # --- AI Assistant Query ---
     return "chat"