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

autosh/config-template.toml

@@ -0,0 +1,13 @@
+# autosh configuration file
+[autosh]
+api_key = "sk-or-v1-..."
+# model = "openai/gpt-4.1"
+# think_model = "openai/o4-mini-high"
+
+[plugins]
+calc = {}
+cli = {}
+clock = {}
+code = {}
+# search = { tavily_api_key = "tvly-dev-..." }
+# web = { tavily_api_key = "tvly-dev-..." }

autosh/config.py

@@ -3,6 +3,8 @@ from pydantic import BaseModel, Field
 from pathlib import Path
 import tomllib
 
+import rich
+
 USER_CONFIG_PATH = Path.home() / ".config" / "autosh" / "config.toml"
 
 
@@ -41,11 +43,20 @@ class Config(BaseModel):
 
     @staticmethod
     def load() -> "Config":
+        if not USER_CONFIG_PATH.is_file():
+            # Copy config.template.toml to USER_CONFIG_PATH
+            USER_CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
+            template = Path(__file__).parent / "config.template.toml"
+            USER_CONFIG_PATH.write_text(template.read_text())
         if USER_CONFIG_PATH.is_file():
-            doc = tomllib.loads(USER_CONFIG_PATH.read_text())
-            main = doc.get("autosh", {})
-            plugins = Plugins(**doc.get("plugins", {}))
-            config = Config.model_validate({**main, "plugins": plugins})
+            try:
+                doc = tomllib.loads(USER_CONFIG_PATH.read_text())
+                main = doc.get("autosh", {})
+                plugins = Plugins(**doc.get("plugins", {}))
+                config = Config.model_validate({**main, "plugins": plugins})
+            except tomllib.TOMLDecodeError as e:
+                rich.print(f"[bold red]Error:[/bold red] invalid config file: {e}")
+                sys.exit(1)
         else:
             config = Config()
         return config

autosh/main.py

@@ -8,7 +8,7 @@ from rich.columns import Columns
 from rich.panel import Panel
 import argparse
 
-from autosh.config import CLI_OPTIONS, CONFIG
+from autosh.config import CLI_OPTIONS, CONFIG, USER_CONFIG_PATH
 from .session import Session
 import sys
 
@@ -149,14 +149,13 @@ def main():
     # dotenv.load_dotenv()
     prompt, args = parse_args()
 
+    if key := os.getenv("OPENROUTER_API_KEY"):
+        CONFIG.api_key = key
     if CONFIG.api_key is None:
-        if key := os.getenv("OPENROUTER_API_KEY"):
-            CONFIG.api_key = key
-        else:
-            rich.print(
-                "[bold red]Error:[/bold red] [red]No API key found. Please set the OPENROUTER_API_KEY environment variable or add it to your config file.[/red]"
-            )
-            sys.exit(1)
+        rich.print(
+            f"[bold red]Error:[/bold red] [red]OpenRouter API key not found.\nPlease set the OPENROUTER_API_KEY environment variable or add it to your config file: {USER_CONFIG_PATH}.[/red]"
+        )
+        sys.exit(1)
     try:
         asyncio.run(start_session(prompt, args))
     except (KeyboardInterrupt, EOFError):

autosh/md/__init__.py

@@ -0,0 +1,8 @@
+from typing import AsyncGenerator
+
+from autosh.md.printer import StreamedMarkdownPrinter
+
+
+async def stream_md(stream: AsyncGenerator[str, None]):
+    mp = StreamedMarkdownPrinter(stream)
+    await mp.parse_doc()

autosh/md/inline_text.py

@@ -0,0 +1,214 @@
+from dataclasses import dataclass, field
+from autosh.md.printer import StreamedMarkdownPrinter
+from autosh.md.state import State
+
+
+@dataclass
+class Keyword:
+    token: str
+
+    def is_bold(self) -> bool:
+        return self.token[0] in ["*", "_"] and len(self.token) == 2
+
+    def is_italic(self) -> bool:
+        return self.token[0] in ["*", "_"] and len(self.token) == 1
+
+    def is_bold_and_italic(self) -> bool:
+        return self.token[0] in ["*", "_"] and len(self.token) > 2
+
+    def is_bold_or_italic(self) -> bool:
+        return self.token[0] in ["*", "_"]
+
+
+class InvalidState(Exception):
+    def __init__(self, token: str):
+        super().__init__(token)
+        self.token = token
+
+
+@dataclass
+class InlineScope:
+    state: State
+
+    stack: list[tuple[Keyword, State]] = field(default_factory=list)
+
+    @property
+    def last(self) -> str | None:
+        return self.stack[-1][0].token if len(self.stack) > 0 else None
+
+    @property
+    def strike(self) -> bool:
+        return self.last == "~~"
+
+    @property
+    def code(self) -> bool:
+        return self.last == "`"
+
+    def enter(self, kind: Keyword):
+        match kind:
+            case "`":
+                state = self.state._enter_scope(dim=True)
+            case "~~":
+                state = self.state._enter_scope(strike=True)
+            case _ if kind.is_bold():
+                state = self.state._enter_scope(bold=True)
+            case _ if kind.is_italic():
+                state = self.state._enter_scope(italic=True)
+            case _ if kind.is_bold_and_italic():
+                state = self.state._enter_scope(bold=True, italic=True)
+            case _:
+                state = self.state._enter_scope()
+        self.stack.append((kind, state))
+
+    def exit(self):
+        t = self.stack.pop()
+        state = t[1]
+        self.state._exit_scope(state)
+
+
+class InlineTextPrinter:
+    def __init__(self, p: StreamedMarkdownPrinter, table: bool = False):
+        self.p = p
+        self.terminator = ["\n", None]
+        if table:
+            self.terminator.append("|")
+
+    def emit(self, s: str):
+        self.p.emit(s)
+
+    def peek(self):
+        return self.p.peek()
+
+    async def consume(self, n: int = 1):
+        return await self.p.consume(n)
+
+    async def check(self, s: str, eof: bool | None = None) -> bool:
+        return await self.p.check(s, eof)
+
+    async def parse_inline_unformatted(self):
+        # Parse until another "`" or a newline, or EOF
+        while not self.peek() in self.terminator:
+            c = await self.p.consume()
+            self.emit(c)
+
+    async def parse_inline_code(self):
+        self.emit("`")
+        # Parse until another "`" or a newline, or EOF
+        while not await self.check("`"):
+            c = await self.p.consume()
+            if c in self.terminator:
+                return
+            self.emit(c)
+            if c == "\\":
+                c = await self.p.consume()
+                if c in self.terminator:
+                    return
+                self.emit(c)
+        self.emit("`")
+        await self.p.consume()
+
+    async def get_next_token(self) -> str | Keyword | None:
+        c = self.peek()
+        if c in self.terminator:
+            return None
+        if c == "`":
+            await self.consume()
+            return Keyword("`")
+        if c == "*":
+            s = ""
+            while self.peek() == "*":
+                s += c
+                await self.consume()
+            return Keyword(s)
+        if c == "_":
+            s = ""
+            while self.peek() == "_":
+                s += c
+                await self.consume()
+            return Keyword(s)
+        if c == "~" and await self.check("~~"):
+            s = ""
+            while self.peek() == "~":
+                s += c
+                await self.consume()
+            return Keyword(s)
+        if c == " " or c == "\t":
+            s = ""
+            while self.peek() == " " or self.peek() == "\t":
+                s += c
+                await self.consume()
+            return s
+        s = c or ""
+        if c == "\\":
+            await self.consume()
+            c = self.peek()
+            if c in self.terminator:
+                return None if len(s) == 0 else s
+            s += c or ""
+            await self.consume()
+        else:
+            await self.consume()
+        if len(s) == 0:
+            return None
+        return s
+
+    async def parse_inline(self, consume_trailing_newline: bool = True):
+        last_is_space = False
+        start = True
+        scope = InlineScope(self.p.state)
+
+        while True:
+            t = await self.get_next_token()
+            curr_is_space = False
+
+            if t is None:
+                if consume_trailing_newline:
+                    if self.peek() == "\n":
+                        await self.consume()
+                    self.emit("\n")
+                return
+            elif isinstance(t, str):
+                # Space or normal text
+                if t[0] == " ":
+                    curr_is_space = True
+                    self.emit(" ")
+                else:
+                    self.emit(t)
+            elif t.token == "`":
+                # Inline code
+                scope.enter(t)
+                with self.p.state.style(dim=True):
+                    await self.parse_inline_code()
+                scope.exit()
+            elif t.token == "~~":
+                # Strike through
+                if not scope.strike:
+                    scope.enter(t)
+                    self.emit("~~")
+                else:
+                    self.emit("~~")
+                    scope.exit()
+            elif (
+                t.is_bold_or_italic()
+                and (last_is_space or start)
+                and scope.last != t.token
+            ):
+                # Start bold or italics
+                scope.enter(t)
+                self.emit(t.token)
+            elif (
+                t.is_bold_or_italic()
+                and self.peek() in [" ", "\t", *self.terminator]
+                and scope.last == t.token
+            ):
+                # End bold or italics
+                self.emit(t.token)
+                scope.exit()
+            else:
+                # print(
+                #     "Invalid token:", t.token, f"[{self.peek()}]", scope.last == t.token
+                # )
+                raise InvalidState(t.token)
+
+            last_is_space = curr_is_space
+            start = False

autosh/md/printer.py

@@ -0,0 +1,301 @@
+import os
+from typing import AsyncGenerator, Literal
+import unicodedata
+
+from autosh.md.state import Color, State
+from autosh.md.stream import TextStream
+
+
+class StreamedMarkdownPrinter:
+    def __init__(self, gen: AsyncGenerator[str, None]):
+        self.stream = TextStream(gen)
+        self.state = State()
+
+    def emit(self, s: str):
+        self.state.emit(s)
+
+    def peek(self):
+        return self.stream.peek()
+
+    async def consume(self, n: int = 1):
+        return await self.stream.consume(n)
+
+    async def check(self, s: str, eof: bool | None = None) -> bool:
+        return await self.stream.check(s, eof)
+
+    async def parse_inline(self, consume_trailing_newline: bool = True):
+        from autosh.md.inline_text import InlineTextPrinter
+
+        itp = InlineTextPrinter(self)
+        await itp.parse_inline(consume_trailing_newline)
+
+    async def parse_heading(self):
+        hashes = 0
+        while True:
+            c = await self.consume()
+            if c == "#":
+                hashes += 1
+            else:
+                break
+        match hashes:
+            case 1:
+                with self.state.style(bold=True, bg=Color.MAGENTA):
+                    with self.state.style(dim=True):
+                        self.emit("#" * hashes + " ")
+                    await self.parse_inline(consume_trailing_newline=False)
+            case 2:
+                with self.state.style(bold=True, underline=True, color=Color.MAGENTA):
+                    with self.state.style(dim=True):
+                        self.emit("#" * hashes + " ")
+                    await self.parse_inline(consume_trailing_newline=False)
+            case 3:
+                with self.state.style(bold=True, color=Color.MAGENTA):
+                    with self.state.style(dim=True):
+                        self.emit("#" * hashes + " ")
+                    await self.parse_inline(consume_trailing_newline=False)
+            case 4:
+                with self.state.style(bold=True, italic=True, color=Color.MAGENTA):
+                    with self.state.style(dim=True):
+                        self.emit("#" * hashes + " ")
+                    await self.parse_inline(consume_trailing_newline=False)
+            case _:
+                with self.state.style(bold=True):
+                    with self.state.style(dim=True):
+                        self.emit("#" * hashes + " ")
+                    await self.parse_inline(consume_trailing_newline=False)
+        await self.consume()  # consume the newline
+        self.emit("\n")
+
+    async def parse_paragraph(self):
+        while True:
+            await self.parse_inline()
+            if (
+                self.peek() != "\n"
+                and not await self.stream.non_paragraph_block_start()
+            ):
+                if self.peek() == "\n":
+                    await self.consume()
+                break
+            else:
+                break
+
+    async def parse_multiline_code(self):
+        with self.state.style(dim=True):
+            self.emit("```")
+            await self.consume(3)
+            while not await self.check("\n```"):
+                c = await self.consume()
+                if c is None:
+                    self.emit("\n")
+                    return
+                self.emit(c)
+            self.emit("\n```\n")
+            await self.consume(4)
+
+    async def parse_list(self, ordered: bool):
+        indents = [0]
+        counter = [1]
+        # first item
+        if ordered:
+            self.emit("1. ")
+            await self.consume()
+        else:
+            self.emit("• ")
+        await self.consume()
+        await self.parse_inline()
+        while True:
+            indent = 0
+            while self.peek() in [" ", "\t", "\n"]:
+                if self.peek() in [" ", "\t"]:
+                    indent += 1
+                if self.peek() == "\n":
+                    indent = 0
+                await self.consume()
+            if self.peek() is None:
+                return
+            if ordered and not await self.stream.ordered_list_label():
+                return
+            if not ordered and not await self.stream.unordered_list_label():
+                return
+            if not ordered:
+                await self.consume()
+            else:
+                while self.peek() is not None and self.peek() != ".":
+                    await self.consume()
+                await self.consume()
+
+            depth = None
+            for i in range(len(indents) - 1):
+                if indents[i] <= indent and indents[i + 1] > indent:
+                    depth = i
+                    break
+            if depth is None and indents[-1] + 2 <= indent:
+                # indent one more level
+                indents.append(indent)
+                depth = len(indents) - 1
+                counter.append(1)
+            elif depth is None:
+                # same as last level
+                depth = len(indents) - 1
+                counter[depth] += 1
+            else:
+                # dedent
+                indents = indents[: depth + 1]
+                counter = counter[: depth + 1]
+                counter[depth] += 1
+            if not ordered:
+                self.emit("  " * depth + "• ")
+            else:
+                self.emit("   " * depth + str(counter[depth]) + ". ")
+            await self.parse_inline()
+
+    async def parse_blockquote(self):
+        from autosh.md.inline_text import InlineTextPrinter
+
+        while True:
+            while self.peek() in [" ", "\t"]:
+                await self.consume()
+            if self.peek() != ">":
+                break
+            await self.consume()
+            with self.state.style(bold=True, dim=True):
+                self.emit(">")
+            with self.state.style(dim=True):
+                itp = InlineTextPrinter(self)
+                await itp.parse_inline_unformatted()
+            if self.peek() == "\n":
+                self.emit("\n")
+                await self.consume()
+
+    async def parse_table(self):
+        def str_size(s: str) -> int:
+            size = 0
+            for c in s:
+                match unicodedata.east_asian_width(c):
+                    case "F" | "W":
+                        size += 2
+                    case _:
+                        size += 1
+            return size
+
+        rows: list[list[str]] = []
+        while True:
+            if self.peek() != "|":
+                break
+            s = "|"
+            await self.consume()
+            while self.peek() not in ["\n", None]:
+                s += self.peek() or ""
+                await self.consume()
+            if self.peek() == "\n":
+                await self.consume()
+            s = s.strip("|")
+            rows.append([c.strip() for c in s.split("|")])
+        # not enough rows
+        if len(rows) < 2:
+            self.emit("| " + " | ".join(rows[0]) + " |\n")
+        # do some simple formatting
+        cols = len(rows[0])
+        col_widths = [0] * cols
+        aligns: list[Literal["left", "right", "center"]] = ["left"] * cols
+        for i, row in enumerate(rows):
+            if i == 1:
+                # check for alignment
+                for j, c in enumerate(row):
+                    if c.startswith(":") and c.endswith(":"):
+                        aligns[j] = "center"
+                    elif c.startswith(":"):
+                        aligns[j] = "left"
+                    elif c.endswith(":"):
+                        aligns[j] = "right"
+                continue
+            for j, c in enumerate(row):
+                col_widths[j] = max(col_widths[j], str_size(c))
+        # print top border
+        with self.state.style(dim=True):
+            for j, c in enumerate(rows[0]):
+                self.emit("┌" if j == 0 else "┬")
+                self.emit("─" * (col_widths[j] + 2))
+            self.emit("┐\n")
+        # print the table
+        for i, row in enumerate(rows):
+            for j, c in enumerate(row):
+                with self.state.style(dim=True):
+                    self.emit("│" if i != 1 else ("├" if j == 0 else "┼"))
+                if i == 1:
+                    text = "─" * (col_widths[j] + 2)
+                else:
+                    s = str_size(c)
+                    align = aligns[j] if i != 0 else "center"
+                    match align:
+                        case "left":
+                            text = c + " " * (col_widths[j] - s)
+                        case "right":
+                            text = " " * (col_widths[j] - s) + c
+                        case "center":
+                            padding = col_widths[j] - s
+                            text = " " * (padding // 2) + c
+                            text += " " * (padding - padding // 2)
+                    text = f" {text} "
+                with self.state.style(dim=i == 1):
+                    self.emit(text)
+            with self.state.style(dim=True):
+                self.emit("│\n" if i != 1 else "┤\n")
+        # print bottom border
+        with self.state.style(dim=True):
+            for j, c in enumerate(rows[0]):
+                self.emit("└" if j == 0 else "┴")
+                self.emit("─" * (col_widths[j] + 2))
+            self.emit("┘\n")
+
+    async def parse_doc(self):
+        await self.stream.init()
+        start = True
+        while True:
+            # Remove leading spaces and empty lines
+            indent = 0
+            while self.peek() in [" ", "\t", "\n"]:
+                if self.peek() in [" ", "\t"]:
+                    indent += 1
+                if self.peek() == "\n":
+                    indent = 0
+                await self.consume()
+            if self.peek() is None:
+                break
+            if not start:
+                self.emit("\n")
+            start = False
+            match c := self.peek():
+                case None:
+                    break
+                # Heading
+                case "#":
+                    await self.parse_heading()
+                # Code
+                case "`" if await self.check("```"):
+                    await self.parse_multiline_code()
+                # Separator
+                case _ if await self.check("---"):
+                    await self.consume(3)
+                    width = min(os.get_terminal_size().columns, 80)
+                    with self.state.style(dim=True):
+                        self.emit("─" * width)
+                # Unordered list
+                case _ if await self.stream.unordered_list_label():
+                    await self.parse_list(False)
+                # Ordered list
+                case _ if await self.stream.ordered_list_label():
+                    await self.parse_list(True)
+                # Blockquote
+                case ">":
+                    await self.parse_blockquote()
+                # Table
+                case "|":
+                    await self.parse_table()
+                # Normal paragraph
+                case _:
+                    await self.parse_paragraph()
+        self.emit("\x1b[0m")  # Reset all
+
+    def __await__(self):
+        return self.parse_doc().__await__()