autosh 0.0.0__py3-none-any.whl → 0.0.1__py3-none-any.whl

autosh/plugins/__init__.py

@@ -0,0 +1,87 @@
+import rich
+from rich.prompt import Confirm
+from rich.panel import Panel
+from rich.console import RenderableType
+from autosh.config import CLI_OPTIONS, CONFIG
+
+
+def banner(tag: str, text: str | None = None, dim: str | None = None):
+    if CLI_OPTIONS.quiet:
+        return
+    s = f"[bold magenta]{tag}[/bold magenta]"
+    if text:
+        s += f" [italic magenta]{text}[/italic magenta]"
+    if dim:
+        s += f" [italic dim]{dim}[/italic dim]"
+    s += "\n"
+    rich.print(s)
+
+
+def confirm(message: str):
+    if CLI_OPTIONS.yes:
+        return True
+    result = Confirm.ask(
+        f"[magenta]{message}[/magenta]", default=True, case_sensitive=False
+    )
+    if not CLI_OPTIONS.quiet:
+        rich.print()
+    return result
+
+
+def cmd_preview_panel(title: str, content: RenderableType, short: str | None = None):
+    if CLI_OPTIONS.quiet and not CLI_OPTIONS.yes:
+        if short:
+            rich.print(f"[magenta]{short}[/magenta]\n")
+        return
+    panel = Panel.fit(content, title=f"[magenta]{title}[/magenta]", title_align="left")
+    rich.print(panel)
+    rich.print()
+
+
+def cmd_result_panel(
+    title: str,
+    out: str | None = None,
+    err: str | None = None,
+):
+    if CLI_OPTIONS.quiet:
+        return
+    if isinstance(out, str):
+        out = out.strip()
+    if isinstance(err, str):
+        err = err.strip()
+    if not out and not err:
+        rich.print(title)
+    else:
+        text = out if out else ""
+        text += (("\n---\n" if out else "") + err) if err else ""
+        panel = Panel.fit(text, title=title, title_align="left", style="dim")
+        rich.print(panel)
+    if not CLI_OPTIONS.quiet:
+        rich.print()
+
+
+from . import calc
+from . import clock
+from . import code
+from . import search
+from . import web
+from . import cli
+
+
+def create_plugins():
+    """Get all plugins in the autosh.plugins module."""
+    cfgs = CONFIG.plugins
+    plugins = []
+    if cfgs.calc is not None:
+        plugins.append(calc.CalculatorPlugin(cfgs.calc.model_dump()))
+    if cfgs.cli is not None:
+        plugins.append(cli.CLIPlugin(cfgs.cli.model_dump()))
+    if cfgs.clock is not None:
+        plugins.append(clock.ClockPlugin(cfgs.clock.model_dump()))
+    if cfgs.code is not None:
+        plugins.append(code.CodePlugin(cfgs.code.model_dump()))
+    if cfgs.search is not None:
+        plugins.append(search.SearchPlugin(cfgs.search.model_dump()))
+    if cfgs.web is not None:
+        plugins.append(web.WebPlugin(cfgs.web.model_dump()))
+    return plugins

autosh/plugins/calc.py

@@ -0,0 +1,22 @@
+from agentia.plugins import tool, Plugin
+from typing import Annotated
+from . import banner
+
+
+class CalculatorPlugin(Plugin):
+    NAME = "calc"
+
+    @tool
+    def evaluate(
+        self,
+        expression: Annotated[
+            str, "The math expression to evaluate. Must be an valid python expression."
+        ],
+    ):
+        """
+        Execute a math expression and return the result. The expression must be an valid python expression that can be execuated by `eval()`.
+        """
+        banner("CALC", expression)
+
+        result = eval(expression)
+        return result

autosh/plugins/cli.py

@@ -0,0 +1,222 @@
+import os
+import sys
+from typing import Annotated
+from agentia.plugins import Plugin, tool
+import rich
+import subprocess
+from enum import StrEnum
+
+from autosh.config import CLI_OPTIONS
+
+from . import banner, confirm, cmd_result_panel, cmd_preview_panel
+
+
+class Color(StrEnum):
+    black = "black"
+    red = "red"
+    green = "green"
+    yellow = "yellow"
+    blue = "blue"
+    magenta = "magenta"
+    cyan = "cyan"
+    white = "white"
+    bright_black = "bright_black"
+    bright_red = "bright_red"
+    bright_green = "bright_green"
+    bright_yellow = "bright_yellow"
+    bright_blue = "bright_blue"
+    bright_magenta = "bright_magenta"
+    bright_cyan = "bright_cyan"
+    bright_white = "bright_white"
+    dim = "dim"
+
+
+class CLIPlugin(Plugin):
+    EXIT_CODE = 0
+
+    @tool
+    def print(
+        self,
+        text: Annotated[
+            str,
+            "The text to print. Can be markdown or using python-rich's markup syntax.",
+        ],
+        color: Annotated[Color | None, "The color of the text"] = None,
+        bold: Annotated[bool, "Whether to print the text in bold"] = False,
+        italic: Annotated[bool, "Whether to print the text in italic"] = False,
+        stderr: Annotated[bool, "Whether to print the text to stderr"] = False,
+        end: Annotated[str, "The text to print at the end"] = "\n",
+    ):
+        """
+        Print an important message to the terminal. NOTE: Important message ONLY! Don't use it when you want to say something to the user.
+        """
+        if color:
+            text = f"[{color}]{text}[/{color}]"
+        if bold:
+            text = f"[bold]{text}[/bold]"
+        if italic:
+            text = f"[italic]{text}[/italic]"
+        rich.print(text, file=sys.stderr if stderr else sys.stdout, end=end)
+        return "DONE. You can continue and no need to repeat the text"
+
+    @tool
+    def chdir(self, path: Annotated[str, "The path to the new working directory"]):
+        """
+        Changes the current working directory of the terminal to another directory.
+        """
+        banner("CWD", path)
+        if not os.path.exists(path):
+            raise FileNotFoundError(f"Path `{path}` does not exist.")
+        os.chdir(path)
+
+    @tool
+    def get_argv(self):
+        """
+        Get the command line arguments.
+        """
+        banner("GET ARGV")
+        if not CLI_OPTIONS.script:
+            return CLI_OPTIONS.args
+        return {
+            "script": str(CLI_OPTIONS.script),
+            "args": CLI_OPTIONS.args,
+        }
+
+    @tool
+    def read(
+        self,
+        path: Annotated[str, "The path to the file to read"],
+    ):
+        """
+        Read a file and print its content.
+        """
+        banner("READ", path)
+        if not os.path.exists(path):
+            raise FileNotFoundError(f"File `{path}` does not exist.")
+        if not os.path.isfile(path):
+            raise FileNotFoundError(f"Path `{path}` is not a file.")
+        with open(path, "r") as f:
+            content = f.read()
+        return content
+
+    @tool
+    def write(
+        self,
+        path: Annotated[str, "The path to the file to write"],
+        content: Annotated[str, "The content to write to the file"],
+        create: Annotated[
+            bool, "Whether to create the file if it does not exist"
+        ] = True,
+        append: Annotated[bool, "Whether to append to the file if it exists"] = False,
+    ):
+        """
+        Write or append text content to a file.
+        """
+        banner("WRITE" if not append else "APPEND", path, f"({len(content)} bytes)")
+
+        if not create and not os.path.exists(path):
+            raise FileNotFoundError(f"File `{path}` does not exist.")
+        if not create and not os.path.isfile(path):
+            raise FileNotFoundError(f"Path `{path}` is not a file.")
+        if path == str(CLI_OPTIONS.script):
+            raise FileExistsError(
+                f"No, you cannot overwrite the script file `{path}`. You're likely writing to it by mistake."
+            )
+        if not confirm("Write file?"):
+            return {"error": "The user declined the write operation."}
+        flag = "a" if append else "w"
+        if create:
+            flag += "+"
+        with open(path, flag) as f:
+            f.write(content)
+        return "DONE. You can continue and no need to repeat the text"
+
+    @tool
+    def stdin_readline(
+        self,
+        prompt: Annotated[
+            str | None, "The optional prompt to display before reading from stdin"
+        ] = None,
+    ):
+        """
+        Read a line from stdin.
+        """
+        if not sys.stdin.isatty():
+            raise RuntimeError("stdin is not a terminal.")
+        return input(prompt)
+
+    @tool
+    def stdin_readall(self):
+        """
+        Read all from stdin until EOF.
+        """
+        if sys.stdin.isatty():
+            raise RuntimeError("No piped input. stdin is a terminal.")
+        if not CLI_OPTIONS.quiet:
+            rich.print("[bold magenta]READ ALL STDIN[/bold magenta]\n")
+        if CLI_OPTIONS.stdin_is_script:
+            raise RuntimeError("No piped input from stdin")
+        return sys.stdin.read()
+
+    @tool
+    def exec(
+        self,
+        command: Annotated[
+            str,
+            "The one-liner bash command to execute. This will be directly sent to `bash -c ...` so be careful with the quotes escaping!",
+        ],
+        explanation: Annotated[
+            str,
+            "Explain what this command does, and what are you going to use it for.",
+        ],
+    ):
+        """
+        Run a one-liner bash command
+        """
+
+        def run():
+            cmd = ["bash", "-c", command]
+            return subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+        # Print the command and explanation
+        cmd_preview_panel(
+            title="Run Command",
+            content=f"[magenta][bold]➜[/bold] [italic]{command}[/italic][/magenta]\n\n[dim]{explanation}[/dim]",
+            short=f"[magenta][bold]➜[/bold] [italic]{command}[/italic][/magenta]",
+        )
+
+        # Ask for confirmation
+        if not confirm("Execute this command?"):
+            return {"error": "The user declined to execute the command."}
+
+        # Execute the command
+        proc_result = run()
+
+        # Print the result
+        if not CLI_OPTIONS.quiet:
+            out = proc_result.stdout.decode("utf-8")
+            err = proc_result.stderr.decode("utf-8")
+            if not out and not err:
+                title = "[green][bold]✔[/bold] Command Finished[/green]"
+            else:
+                if proc_result.returncode != 0:
+                    title = f"[red][bold]✘[/bold] Command Failed [{proc_result.returncode}][/red]"
+                else:
+                    title = "[green][bold]✔[/bold] Command Finished[/green]"
+            cmd_result_panel(title, out, err)
+
+        result = {
+            "stdout": proc_result.stdout.decode("utf-8"),
+            "stderr": proc_result.stderr.decode("utf-8"),
+            "returncode": proc_result.returncode,
+            "success": proc_result.returncode == 0,
+        }
+        return result
+
+    @tool
+    def exit(self, exitcode: Annotated[int, "The exit code of this shell session"] = 0):
+        """
+        Exit the current shell session with an optional exit code.
+        """
+        banner("EXIT", str(exitcode))
+        sys.exit(exitcode)

autosh/plugins/clock.py

@@ -0,0 +1,20 @@
+import datetime
+from agentia.plugins import tool, Plugin
+import tzlocal
+
+from autosh.plugins import banner
+
+
+class ClockPlugin(Plugin):
+    @tool
+    def get_current_time(self):
+        """Get the current UTC time in ISO format"""
+        banner("GET TIME")
+        utc = datetime.datetime.now(datetime.timezone.utc).isoformat()
+        local = datetime.datetime.now().isoformat()
+        timezone = tzlocal.get_localzone_name()
+        return {
+            "utc": utc,
+            "local": local,
+            "timezone": timezone,
+        }

autosh/plugins/code.py

@@ -0,0 +1,68 @@
+from agentia.plugins import tool, Plugin
+from typing import Annotated
+import traceback
+from rich.syntax import Syntax
+from rich.console import group
+from contextlib import redirect_stdout, redirect_stderr
+import io
+from . import confirm, cmd_result_panel, cmd_preview_panel
+
+
+class CodePlugin(Plugin):
+    @tool
+    def execute(
+        self,
+        python_code: Annotated[str, "The python code to run."],
+        explanation: Annotated[
+            str, "Explain what this code does, and what are you going to use it for."
+        ],
+    ):
+        """
+        Execute python code and return the result.
+        The python code must be a valid python source file that accepts no inputs.
+        Print results to stdout or stderr.
+        """
+
+        @group()
+        def code_with_explanation():
+            yield Syntax(python_code.strip(), "python")
+            yield "\n[dim]───[/dim]\n"
+            yield f"[dim]{explanation}[/dim]"
+
+        cmd_preview_panel(
+            title="Run Python",
+            content=code_with_explanation(),
+            short=f"[bold]RUN[/bold] [italic]Python Code[/italic]",
+        )
+
+        if not confirm("Execute this code?"):
+            return {"error": "The user declined to execute the command."}
+
+        out = io.StringIO()
+        err = io.StringIO()
+        with redirect_stdout(out):
+            with redirect_stderr(err):
+                try:
+                    exec(python_code, globals())
+                    o = out.getvalue()
+                    e = err.getvalue()
+                    title = "[green][bold]✔[/bold] Finished[/green]"
+                    result = {
+                        "stdout": o,
+                        "stderr": e,
+                        "success": True,
+                    }
+                except Exception as ex:
+                    o = out.getvalue()
+                    e = err.getvalue()
+                    title = "[red][bold]✘[/bold] Failed [/red]"
+                    result = {
+                        "stdout": o,
+                        "stderr": e,
+                        "success": False,
+                        "error": str(ex),
+                        "traceback": repr(traceback.format_exc()),
+                    }
+
+        cmd_result_panel(title, o, e)
+        return result

autosh/plugins/search.py

@@ -0,0 +1,90 @@
+from agentia.plugins import tool, Plugin
+from typing import Annotated, override
+import os
+from tavily import TavilyClient
+
+from autosh.plugins import banner
+
+
+class SearchPlugin(Plugin):
+    @override
+    async def init(self):
+        if "tavily_api_key" in self.config:
+            key = self.config["tavily_api_key"]
+        elif "TAVILY_API_KEY" in os.environ:
+            key = os.environ["TAVILY_API_KEY"]
+        else:
+            raise ValueError("Please set the TAVILY_API_KEY environment variable.")
+        self.__tavily = TavilyClient(api_key=key)
+
+    @tool
+    async def web_search(
+        self,
+        query: Annotated[
+            str, "The search query. Please be as specific and verbose as possible."
+        ],
+    ):
+        """
+        Perform web search on the given query.
+        Returning the top related search results in json format.
+        When necessary, you need to combine this tool with the get_webpage_content tools (if available), to browse the web in depth by jumping through links.
+        """
+        banner("WEB SEARCH", dim=query)
+
+        tavily_results = self.__tavily.search(
+            query=query,
+            search_depth="advanced",
+            # max_results=10,
+            include_answer=True,
+            include_images=True,
+            include_image_descriptions=True,
+        )
+        return tavily_results
+
+    @tool
+    async def news_search(
+        self,
+        query: Annotated[
+            str, "The search query. Please be as specific and verbose as possible."
+        ],
+    ):
+        """
+        Perform news search on the given query.
+        Returning the top related results in json format.
+        """
+        banner("NEWS SEARCH", dim=query)
+
+        tavily_results = self.__tavily.search(
+            query=query,
+            search_depth="advanced",
+            topic="news",
+            # max_results=10,
+            include_answer=True,
+            include_images=True,
+            include_image_descriptions=True,
+        )
+        return tavily_results
+
+    @tool
+    async def finance_search(
+        self,
+        query: Annotated[
+            str, "The search query. Please be as specific and verbose as possible."
+        ],
+    ):
+        """
+        Search for finance-related news and information on the given query.
+        Returning the top related results in json format.
+        """
+        banner("FINANCE SEARCH", dim=query)
+
+        tavily_results = self.__tavily.search(
+            query=query,
+            search_depth="advanced",
+            topic="finance",
+            # max_results=10,
+            include_answer=True,
+            include_images=True,
+            include_image_descriptions=True,
+        )
+        return tavily_results

autosh/plugins/web.py

@@ -0,0 +1,73 @@
+from io import BytesIO
+import os
+from agentia.plugins import tool, Plugin
+from typing import Annotated, override
+import requests
+from markdownify import markdownify
+import uuid
+from tavily import TavilyClient
+
+from autosh.plugins import banner
+
+
+class WebPlugin(Plugin):
+
+    @override
+    async def init(self):
+        if "tavily_api_key" in self.config:
+            key = self.config["tavily_api_key"]
+        elif "TAVILY_API_KEY" in os.environ:
+            key = os.environ["TAVILY_API_KEY"]
+        else:
+            raise ValueError("Please set the TAVILY_API_KEY environment variable.")
+        self.__tavily = TavilyClient(api_key=key)
+
+    def __embed_file(self, content: bytes, file_ext: str):
+        assert self.agent.knowledge_base is not None
+        with BytesIO(content) as f:
+            ext = file_ext if file_ext.startswith(".") else "." + file_ext
+            f.name = str(uuid.uuid4()) + ext
+            # self.agent.knowledge_base.add_document(f)
+            # file_name = f.name
+            raise NotImplementedError
+        return {
+            "file_name": file_name,
+            "hint": f"This is a .{file_ext} file and it is embeded in the knowledge base. Use _file_search to query the content.",
+        }
+
+    def __get(self, url: str):
+        res = requests.get(url)
+        res.raise_for_status()
+        content_type = res.headers.get("content-type")
+        # if content_type == "application/pdf":
+        #     # Add this file to the knowledge base
+        #     if self.agent.knowledge_base is not None:
+        #         return self.__embed_file(res.content, "pdf")
+        #     return {"content": "This is a PDF file. You don't know how to view it."}
+        md = markdownify(res.text)
+        return {"content": md}
+
+    @tool
+    def get_webpage_content(
+        self,
+        url: Annotated[str, "The URL of the web page to get the content of"],
+    ):
+        """
+        Access a web page by a URL, and fetch the content of this web page (in markdown format).
+        You can always use this tool to directly access web content or access external sites.
+        Use it at any time when you think you may need to access the internet.
+        """
+        banner("BROWSE", dim=url)
+
+        result = self.__tavily.extract(
+            urls=url,
+            # extract_depth="advanced",
+            include_images=True,
+        )
+        failed_results = result.get("failed_results", [])
+        if len(failed_results) > 0:
+            try:
+                return self.__get(url)
+            except Exception as e:
+                pass
+        return result