ara-cli 0.1.10.5__py3-none-any.whl → 0.1.13.3__py3-none-any.whl

ara_cli/chat.py

@@ -1,16 +1,35 @@
 import os
-import errno
+import sys
 import argparse
 import cmd2
 
 from ara_cli.prompt_handler import send_prompt
 
+from . import (
+    CATEGORY_CHAT_CONTROL,
+    CATEGORY_LLM_CONTROL,
+    CATEGORY_SCRIPT_CONTROL,
+    CATEGORY_AGENT_CONTROL,
+)
+from . import ROLE_PROMPT, ROLE_RESPONSE, INTRO
+from . import BINARY_TYPE_MAPPING, DOCUMENT_TYPE_EXTENSIONS
+
 from . import error_handler
 from ara_cli.error_handler import AraError, AraConfigurationError
 
 from ara_cli.file_loaders.document_file_loader import DocumentFileLoader
 from ara_cli.file_loaders.binary_file_loader import BinaryFileLoader
 from ara_cli.file_loaders.text_file_loader import TextFileLoader
+from ara_cli.chat_agent.agent_process_manager import AgentProcessManager
+
+from ara_cli.chat_script_runner.script_runner import ScriptRunner
+from ara_cli.chat_script_runner.script_completer import ScriptCompleter
+from ara_cli.chat_script_runner.script_lister import ScriptLister
+from ara_cli.chat_web_search.web_search import (
+    perform_web_search_completion,
+    is_web_search_supported,
+    get_supported_models_message,
+)
 
 
 extract_parser = argparse.ArgumentParser()
@@ -34,39 +53,6 @@ load_parser.add_argument(
 
 
 class Chat(cmd2.Cmd):
-    CATEGORY_CHAT_CONTROL = "Chat control commands"
-    CATEGORY_LLM_CONTROL = "Language model controls"
-
-    INTRO = """/***************************************/
-                 araarar
-               aa       ara
-             aa    aa   aara
-             a        araarar
-             a        ar  ar
-           aa          ara
-          a               a
-          a               aa
-           a              a
-   ar      aa           aa
-    (c) ara chat by talsen team
-              aa      aa
-               aa    a
-                a aa
-                 aa
-/***************************************/
-Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat mode):"""
-
-    ROLE_PROMPT = "ara prompt"
-    ROLE_RESPONSE = "ara response"
-
-    BINARY_TYPE_MAPPING = {
-        ".png": "image/png",
-        ".jpg": "image/jpeg",
-        ".jpeg": "image/jpeg",
-    }
-
-    DOCUMENT_TYPE_EXTENSIONS = [".docx", ".doc", ".odt", ".pdf"]
-
     def __init__(
         self,
         chat_name: str,
@@ -74,6 +60,7 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
         enable_commands: list[str] | None = None,
     ):
         from ara_cli.template_loader import TemplateLoader
+
         shortcuts = dict(cmd2.DEFAULT_SHORTCUTS)
         if enable_commands:
             enable_commands.append("quit")  # always allow quitting
@@ -97,15 +84,21 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
             self.disable_commands(commands_to_disable)
 
         self.prompt = "ara> "
-        self.intro = Chat.INTRO
+        self.intro = INTRO
 
-        self.default_chat_content = f"# {Chat.ROLE_PROMPT}:\n"
+        self.default_chat_content = f"# {ROLE_PROMPT}:\n"
         self.chat_name = self.setup_chat(chat_name, reset)
         self.chat_name = os.path.abspath(self.chat_name)
         self.chat_history = []
         self.message_buffer = []
         self.config = self._retrieve_ara_config()
         self.template_loader = TemplateLoader(chat_instance=self)
+        self.script_runner = ScriptRunner(chat_instance=self)
+        self.script_lister = ScriptLister()
+        self.script_completer = ScriptCompleter()
+
+        # Initialize agent process manager
+        self.agent_manager = AgentProcessManager(self)
 
     def disable_commands(self, commands: list[str]):
         for command in commands:
@@ -127,6 +120,7 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
         self.aliases["h"] = "help"
         self.aliases["n"] = "NEW"
         self.aliases["e"] = "EXTRACT"
+        self.aliases["SEARCH"] = "search"
         self.aliases["l"] = "LOAD"
         self.aliases["lr"] = "LOAD_RULES"
         self.aliases["li"] = "LOAD_INTENTION"
@@ -134,6 +128,11 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
         self.aliases["lg"] = "LOAD_GIVENS"
         self.aliases["lb"] = "LOAD_BLUEPRINT"
         self.aliases["lt"] = "LOAD_TEMPLATE"
+        self.aliases["rpy"] = "run_pyscript"
+        self.aliases["a"] = "AGENT_RUN"
+        self.aliases["al"] = "LIST_AGENTS"
+        self.aliases["la"] = "LIST_AGENTS"
+        self.aliases["AGENT_LIST"] = "LIST_AGENTS"
 
     def setup_chat(self, chat_name, reset: bool = None):
         if os.path.exists(chat_name):
@@ -148,9 +147,12 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
         chat_file_short = os.path.split(chat_file)[-1]
 
         if reset is None:
-            user_input = input(
-                f"{chat_file_short} already exists. Do you want to reset the chat? (y/N): "
+            print(
+                f"{chat_file_short} already exists. Do you want to reset the chat? (y/N): ",
+                end="",
+                flush=True,
             )
+            user_input = sys.stdin.readline().strip()
             if user_input.lower() == "y":
                 self.create_empty_chat_file(chat_file)
         if reset:
@@ -178,13 +180,14 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
                 print(f"File {file_name} not found.")
                 return False
             return method(self, file_path, *args, **kwargs)
+
         return wrapper
 
     @staticmethod
     def get_last_role_marker(lines):
         if not lines:
             return
-        role_markers = [f"# {Chat.ROLE_PROMPT}:", f"# {Chat.ROLE_RESPONSE}"]
+        role_markers = [f"# {ROLE_PROMPT}:", f"# {ROLE_RESPONSE}"]
         for line in reversed(lines):
             stripped_line = line.strip()
             if stripped_line.startswith(tuple(role_markers)):
@@ -246,8 +249,8 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
         import re
         from ara_cli.prompt_handler import prepend_system_prompt
 
-        prompt_marker = f"# {self.ROLE_PROMPT}:"
-        response_marker = f"# {self.ROLE_RESPONSE}:"
+        prompt_marker = f"# {ROLE_PROMPT}:"
+        response_marker = f"# {ROLE_RESPONSE}:"
 
         split_pattern = re.compile(f"({prompt_marker}|{response_marker})")
 
@@ -283,7 +286,7 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
     def send_message(self):
         self.chat_history = self.load_chat_history(self.chat_name)
         prompt_to_send = self.assemble_prompt()
-        role_marker = f"# {Chat.ROLE_RESPONSE}:"
+        role_marker = f"# {ROLE_RESPONSE}:"
 
         with open(self.chat_name, "a+", encoding="utf-8") as file:
             last_line = self.get_last_line(file)
@@ -332,9 +335,9 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
     def find_last_reply_index(self, lines: list[str]):
         index_to_remove = None
         for i, line in enumerate(reversed(lines)):
-            if line.strip().startswith(f"# {Chat.ROLE_PROMPT}"):
+            if line.strip().startswith(f"# {ROLE_PROMPT}"):
                 break
-            if line.strip().startswith(f"# {Chat.ROLE_RESPONSE}"):
+            if line.strip().startswith(f"# {ROLE_RESPONSE}"):
                 index_to_remove = len(lines) - i - 1
                 break
         return index_to_remove
@@ -359,7 +362,8 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
     def add_prompt_tag_if_needed(self, chat_file: str):
         with open(chat_file, "r", encoding="utf-8") as file:
             lines = file.readlines()
-        prompt_tag = f"# {Chat.ROLE_PROMPT}:"
+
+        prompt_tag = f"# {ROLE_PROMPT}:"
         if Chat.get_last_role_marker(lines) == prompt_tag:
             return
         append = prompt_tag
@@ -369,7 +373,6 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
         with open(chat_file, "a", encoding="utf-8") as file:
             file.write(append)
 
-    # @file_exists_check
     def load_text_file(
         self,
         file_path,
@@ -387,7 +390,6 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
             extract_images=extract_images,
         )
 
-    # @file_exists_check
     def load_binary_file(
         self, file_path, mime_type: str, prefix: str = "", suffix: str = ""
     ):
@@ -401,7 +403,6 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
         reader = MarkdownReader(file_path)
         return reader.read(extract_images=extract_images)
 
-    # @file_exists_check
     def load_document_file(
         self,
         file_path: str,
@@ -427,8 +428,8 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
         block_delimiter: str = "",
         extract_images: bool = False,
     ):
-        binary_type_mapping = Chat.BINARY_TYPE_MAPPING
-        document_type_extensions = Chat.DOCUMENT_TYPE_EXTENSIONS
+        binary_type_mapping = BINARY_TYPE_MAPPING
+        document_type_extensions = DOCUMENT_TYPE_EXTENSIONS
 
         file_type = None
         file_name_lower = file_name.lower()
@@ -467,7 +468,8 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
             files.sort()
             for i, file in enumerate(files):
                 print(f"{i + 1}: {os.path.basename(file)}")
-            choice = input("Please choose a file to load (enter number): ")
+            print("Please choose a file to load (enter number): ", end="", flush=True)
+            choice = sys.stdin.readline().strip()
             try:
                 choice_index = int(choice) - 1
                 if choice_index < 0 or choice_index >= len(files):
@@ -493,6 +495,7 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
 
     def do_quit(self, _):
         """Exit ara-cli"""
+        self.agent_manager.cleanup_agent_process()
         print("Chat ended")
         self.last_result = True
         return True
@@ -562,7 +565,8 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
         import glob
 
         if file_name == "":
-            file_name = input("What file do you want to load? ")
+            print("What file do you want to load? ", end="", flush=True)
+            file_name = sys.stdin.readline().strip()
         file_pattern = os.path.join(os.path.dirname(self.chat_name), file_name)
         matching_files = glob.glob(file_pattern)
         if not matching_files:
@@ -573,7 +577,7 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
         return matching_files
 
     def load_image(self, file_name: str, prefix: str = "", suffix: str = ""):
-        binary_type_mapping = Chat.BINARY_TYPE_MAPPING
+        binary_type_mapping = BINARY_TYPE_MAPPING
 
         file_type = None
         file_name_lower = file_name.lower()
@@ -602,6 +606,62 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
             return False
         return True
 
+    @cmd2.with_category(CATEGORY_CHAT_CONTROL)
+    def do_search(self, query: str):
+        """Perform a web search and append the results to the chat.
+        Usage: search <query>
+        """
+        if not query:
+            self.poutput("Please provide a search query.")
+            return
+
+        # Check if web search is supported by the current model
+        from ara_cli.prompt_handler import LLMSingleton
+
+        chat_instance = LLMSingleton.get_instance()
+        config_parameters = chat_instance.get_config_by_purpose("default")
+        default_llm = config_parameters.get("model")
+
+        is_supported, _ = is_web_search_supported(default_llm)
+        if not is_supported:
+            self.poutput(get_supported_models_message(default_llm))
+            return
+
+        self.add_prompt_tag_if_needed(self.chat_name)
+
+        role_marker = f"# Web Search Results for '{query}':"
+
+        with open(self.chat_name, "a+", encoding="utf-8") as file:
+            last_line = self.get_last_line(file)
+
+            self.poutput(role_marker)
+
+            if not last_line.startswith(role_marker):
+                if last_line:
+                    file.write("\n")
+                file.write(role_marker + "\n")
+
+            try:
+                # perform_web_search_completion now returns a generator or a string
+                search_result = perform_web_search_completion(query)
+
+                if isinstance(search_result, str):
+                    # If it's a string, it's an error/info message
+                    self.poutput(search_result)
+                    file.write(search_result + "\n")
+                else:
+                    # Otherwise, it's a generator, stream the content
+                    for chunk in search_result:
+                        chunk_content = chunk.choices[0].delta.content
+                        if not chunk_content:
+                            continue
+                        self.poutput(chunk_content, end="")
+                        file.write(chunk_content)
+                        file.flush()
+                    self.poutput("")
+            except Exception as e:
+                error_handler.report_error(e)
+
     @cmd2.with_category(CATEGORY_CHAT_CONTROL)
     def do_LOAD_IMAGE(self, file_name):
         """Load an image file and append it to chat file. Can be given the file name in-line. Will attempt to find the file relative to chat file first, then treat the given path as absolute"""
@@ -619,7 +679,7 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
                 # Determine mime type
                 file_type = None
                 file_path_lower = file_path.lower()
-                for extension, mime_type in Chat.BINARY_TYPE_MAPPING.items():
+                for extension, mime_type in BINARY_TYPE_MAPPING.items():
                     if file_path_lower.endswith(extension):
                         file_type = mime_type
                         break
@@ -726,7 +786,8 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
     def do_NEW(self, chat_name):
         """Create a new chat. Optionally provide a chat name in-line: NEW new_chat"""
         if chat_name == "":
-            chat_name = input("What should be the new chat name? ")
+            print("What should be the new chat name? ", end="", flush=True)
+            chat_name = sys.stdin.readline().strip()
         current_directory = os.path.dirname(self.chat_name)
         chat_file_path = os.path.join(current_directory, chat_name)
         self.__init__(chat_file_path)
@@ -739,7 +800,8 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
     @cmd2.with_category(CATEGORY_CHAT_CONTROL)
     def do_CLEAR(self, _):
         """Clear the chat and the file containing it"""
-        user_input = input("Are you sure you want to clear the chat? (y/N): ")
+        print("Are you sure you want to clear the chat? (y/N): ", end="", flush=True)
+        user_input = sys.stdin.readline().strip()
         if user_input.lower() != "y":
             return
         self.create_empty_chat_file(self.chat_name)
@@ -750,17 +812,23 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
     @cmd2.with_category(CATEGORY_CHAT_CONTROL)
     def do_LOAD_RULES(self, rules_name):
         """Load rules from ./prompt.data/*.rules.md or from a specified template directory if an argument is given. Specify global/<rules_template> to access globally defined rules templates"""
-        self.template_loader.load_template(rules_name, "rules", self.chat_name, "*.rules.md")
+        self.template_loader.load_template(
+            rules_name, "rules", self.chat_name, "*.rules.md"
+        )
 
     @cmd2.with_category(CATEGORY_CHAT_CONTROL)
     def do_LOAD_INTENTION(self, intention_name):
         """Load intention from ./prompt.data/*.intention.md or from a specified template directory if an argument is given. Specify global/<intention_template> to access globally defined intention templates"""
-        self.template_loader.load_template(intention_name, "intention", self.chat_name, "*.intention.md")
+        self.template_loader.load_template(
+            intention_name, "intention", self.chat_name, "*.intention.md"
+        )
 
     @cmd2.with_category(CATEGORY_CHAT_CONTROL)
     def do_LOAD_COMMANDS(self, commands_name):
         """Load commands from ./prompt.data/*.commands.md or from a specified template directory if an argument is given. Specify global/<commands_template> to access globally defined commands templates"""
-        self.template_loader.load_template(commands_name, "commands", self.chat_name, "*.commands.md")
+        self.template_loader.load_template(
+            commands_name, "commands", self.chat_name, "*.commands.md"
+        )
 
     @cmd2.with_category(CATEGORY_CHAT_CONTROL)
     def do_LOAD_BLUEPRINT(self, blueprint_name):
@@ -934,7 +1002,7 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
     def do_SEND(self, _):
         """Send prompt to the LLM"""
         message = "\n".join(self.message_buffer)
-        self.save_message(Chat.ROLE_PROMPT, message)
+        self.save_message(ROLE_PROMPT, message)
         self.send_message()
 
     @cmd2.with_category(CATEGORY_CHAT_CONTROL)
@@ -1053,7 +1121,9 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
 
     def _template_completer(self, text: str, template_type: str) -> list[str]:
         """Generic completer for different template types."""
-        available_templates = self.template_loader.get_available_templates(template_type, os.path.dirname(self.chat_name))
+        available_templates = self.template_loader.get_available_templates(
+            template_type, os.path.dirname(self.chat_name)
+        )
         if not text:
             return available_templates
         return [t for t in available_templates if t.startswith(text)]
@@ -1072,4 +1142,155 @@ Start chatting (type 'HELP'/'h' for available commands, 'QUIT'/'q' to exit chat
 
     def complete_LOAD_BLUEPRINT(self, text, line, begidx, endidx):
         """Completer for the LOAD_BLUEPRINT command."""
-        return self._template_completer(text, "blueprint")
+        return self._template_completer(text, "blueprint")
+
+    def _select_script_from_list(
+        self, scripts: list[str], not_found_message: str, prompt: str
+    ) -> str | None:
+        """Displays a list of scripts and prompts the user to select one."""
+        if not scripts:
+            self.poutput(not_found_message)
+            return None
+
+        # Sort the scripts alphabetically by their basename for consistent display
+        # Create a list of (basename, full_script_name) tuples for sorting and later retrieval
+        scripts_with_basenames = [(os.path.basename(s), s) for s in scripts]
+        scripts_with_basenames.sort(
+            key=lambda x: x[0].lower()
+        )  # Sort by lowercase basename
+
+        for i, (basename, full_script_name) in enumerate(scripts_with_basenames):
+            self.poutput(f"{i + 1}: {basename}")
+
+        try:
+            choice = input(prompt)
+            choice_index = int(choice) - 1
+            if 0 <= choice_index < len(scripts_with_basenames):
+                # Return the full script name from the sorted list
+                return scripts_with_basenames[choice_index][1]
+            else:
+                self.poutput("Invalid choice. Aborting.")
+                return None
+        except (ValueError, EOFError):
+            self.poutput("Invalid input. Aborting.")
+            return None
+
+    @cmd2.with_category(CATEGORY_SCRIPT_CONTROL)
+    def do_run_pyscript(self, args):
+        """Run a python script from the chat.
+        Usage: run_pyscript <script_name> [args...]
+        """
+        script_name, script_args = self._parse_run_pyscript_args(args)
+
+        # If no script name provided, list available scripts grouped by type
+        if not script_name:
+            self._list_available_scripts()
+            return
+
+        script_to_run = self._resolve_script_to_run(script_name, script_args)
+
+        if not script_to_run:
+            return
+
+        # Pass arguments to script runner
+        output = self.script_runner.run_script(script_to_run, script_args)
+        if output:
+            self.poutput(output.strip())
+
+    def _list_available_scripts(self):
+        """Lists available scripts grouped by type (global and custom)."""
+        global_scripts = self.script_lister.get_global_scripts()
+        custom_scripts = self.script_lister.get_custom_scripts()
+
+        if not global_scripts and not custom_scripts:
+            self.poutput("No scripts found.")
+            return
+
+        self.poutput("Available scripts:")
+        self.poutput("")
+
+        if custom_scripts:
+            self.poutput("Custom scripts:")
+            for script in sorted(custom_scripts):
+                self.poutput(f"  {script}")
+            self.poutput("")
+
+        if global_scripts:
+            self.poutput("Global scripts:")
+            for script in sorted(global_scripts):
+                self.poutput(f"  global/{script}")
+
+    def _parse_run_pyscript_args(self, args):
+        """Parses arguments for run_pyscript command."""
+        import shlex
+
+        if not args:
+            return "", []
+
+        # args is a cmd2.Statement (subclass of str), so we can use it directly
+        full_args = str(args)
+        # Use shlex to split arguments, enabling quoted args support
+        split_args = shlex.split(full_args)
+        if not split_args:
+            return "", []
+
+        script_name = split_args[0]
+        script_args = split_args[1:] if len(split_args) > 1 else []
+        return script_name, script_args
+
+    def _resolve_script_to_run(self, script_name, script_args):
+        """Resolves the script name to run."""
+        return script_name
+
+    def complete_run_pyscript(self, text, line, begidx, endidx):
+        """Completer for the run_pyscript command."""
+        # Get all scripts: ['custom.py', 'global/global.py']
+        available_scripts = self.script_lister.get_all_scripts()
+
+        # Add special commands
+        special_commands = [
+            # "global/"
+            # "*", "global/*"
+        ]
+
+        possible_completions = sorted(list(set(available_scripts + special_commands)))
+
+        # Filter based on what the user has typed
+        return [s for s in possible_completions if s.startswith(text)]
+
+    # ===== AGENT CONTROL COMMANDS =====
+
+    @cmd2.with_category(CATEGORY_AGENT_CONTROL)
+    def do_LIST_AGENTS(self, _):
+        """Lists all available executable binary agents."""
+        from ara_cli.commands.list_agents_command import ListAgentsCommand
+
+        command = ListAgentsCommand(chat_instance=self)
+        command.execute()
+
+    @cmd2.with_category(CATEGORY_AGENT_CONTROL)
+    def do_AGENT_RUN(self, args):
+        """Run a binary agent interactively from the 'ara/.araconfig/agents' directory.
+        Usage: AGENT_RUN <agent_name> [arg1] [arg2] ...
+        Example:
+          AGENT_RUN feature-creation -b .
+        """
+
+        from ara_cli.commands.agent_run_command import AgentRunCommand
+
+        command = AgentRunCommand(self, args)
+        command.execute()
+
+    def complete_AGENT_RUN(self, text, line, begidx, endidx):
+        """Completer for AGENT_RUN command."""
+        from ara_cli.commands.list_agents_command import list_available_binary_agents
+
+        parts = line.split()
+        # This completer runs when the user is typing the first argument (the agent name)
+        if len(parts) < 2 or (len(parts) == 2 and not line.endswith(" ")):
+            available_agents = list_available_binary_agents(self)
+            if not text:
+                return available_agents
+            return [a for a in available_agents if a.startswith(text)]
+        # For subsequent arguments, we can offer file/directory completion
+        return self.path_complete(text, line, begidx, endidx)