simonw/llm + python rich

stream_3.py

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.12"
# dependencies = ["rich", "aiofiles"]
# ///

# version 3: aiofiles

import asyncio
import sys
import os
from rich.console import Console, ConsoleOptions, RenderResult
from rich.live import Live
from rich.markdown import Markdown, CodeBlock
from rich.syntax import Syntax
from rich.text import Text
import aiofiles

async def read_stdin_async():
    """Async generator that yields chunks from stdin"""
    loop = asyncio.get_event_loop()
    reader = asyncio.StreamReader()
    protocol = asyncio.StreamReaderProtocol(reader)
    
    await loop.connect_read_pipe(lambda: protocol, sys.stdin)
    
    while True:
        chunk = await reader.read(1024)
        if not chunk:
            break
        yield chunk.decode('utf-8', errors='replace')


async def main():
    prettier_code_blocks()
    console = Console()
    
    # Check if stdin has data
    if sys.stdin.isatty():
        console.print("[red]No input received from stdin.[/red]")
        console.print("[dim]Example: echo '# Hello World' | ./stream[/dim]")
        return
    
    buffer = ""
    
    with Live("", console=console, vertical_overflow="visible", refresh_per_second=10) as live:
        try:
            async for chunk in read_stdin_async():
                buffer += chunk
                live.update(Markdown(buffer))
                
        except KeyboardInterrupt:
            console.print("\n[yellow]Interrupted by user[/yellow]")
        except Exception as e:
            console.print(f"\n[red]Error: {e}[/red]")
        finally:
            # Ensure final buffer is displayed
            if buffer and not live.is_started:
                console.print(Markdown(buffer))


def prettier_code_blocks():
    """Make rich code blocks prettier and easier to copy.

    From https://github.com/samuelcolvin/aicli/blob/v0.8.0/samuelcolvin_aicli.py#L22
    """

    class SimpleCodeBlock(CodeBlock):
        def __rich_console__(
            self, console: Console, options: ConsoleOptions
        ) -> RenderResult:
            code = str(self.text).rstrip()
            yield Text(self.lexer_name, style='dim')
            yield Syntax(
                code,
                self.lexer_name,
                theme=self.theme,
                background_color='default',
                word_wrap=True,
            )
            yield Text(f'/{self.lexer_name}', style='dim')

    Markdown.elements['fence'] = SimpleCodeBlock


if __name__ == '__main__':
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass

# vim: set ft=py:

stream_v1.py

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.12"
# dependencies = ["rich"]
# ///

# version 1: inspired from pydantic-ai, fake simulated stream

import asyncio
import sys
from rich.console import Console, ConsoleOptions, RenderResult
from rich.live import Live
from rich.markdown import Markdown, CodeBlock
from rich.syntax import Syntax
from rich.text import Text

async def main():
    prettier_code_blocks()
    console = Console()

    # Read from stdin
    console.print("\nResponse:", style='cyan')

    # Check if stdin is a pipe/has data
    if not sys.stdin.isatty():
        # Read all input from stdin
        input_text = sys.stdin.read()

        # Display the input using Rich's Live display with Markdown rendering
        with Live('', console=console, vertical_overflow='visible') as live:
            # If you want to simulate streaming, you can do it character by character
            # or line by line. Here's a character-by-character approach:
            accumulated_text = ""
            for char in input_text:
                accumulated_text += char
                live.update(Markdown(accumulated_text))
                await asyncio.sleep(0.001)  # Small delay to simulate streaming

        # Alternatively, if you just want to display it all at once:
        # console.print(Markdown(input_text))
    else:
        console.log("No input received from stdin. Pipe some text to this script.", style='red')
        console.log("Example: echo '# Hello World' | python script.py", style='dim')


def prettier_code_blocks():
    """Make rich code blocks prettier and easier to copy.

    From https://github.com/samuelcolvin/aicli/blob/v0.8.0/samuelcolvin_aicli.py#L22
    """

    class SimpleCodeBlock(CodeBlock):
        def __rich_console__(
            self, console: Console, options: ConsoleOptions
        ) -> RenderResult:
            code = str(self.text).rstrip()
            yield Text(self.lexer_name, style='dim')
            yield Syntax(
                code,
                self.lexer_name,
                theme=self.theme,
                background_color='default',
                word_wrap=True,
            )
            yield Text(f'/{self.lexer_name}', style='dim')

    Markdown.elements['fence'] = SimpleCodeBlock


if __name__ == '__main__':
    asyncio.run(main())

# vim: set ft=py:

stream_v2.py

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.12"
# dependencies = ["rich"]
# ///

# version 2: chunk size to control stream

import asyncio
import sys
import select
from rich.console import Console, ConsoleOptions, RenderResult
from rich.live import Live
from rich.markdown import Markdown, CodeBlock
from rich.syntax import Syntax
from rich.text import Text

async def main():
    prettier_code_blocks()
    console = Console()

    # Check if stdin has data
    if sys.stdin.isatty():
        console.print("[red]No input received from stdin.[/red]")
        console.print("[dim]Example: echo '# Hello World' | ./stream[/dim]")
        return

    buffer = ""
    chunk_size = 30  # Read in larger chunks for better performance

    with Live("", console=console, vertical_overflow="visible", refresh_per_second=10) as live:
        try:
            while True:
                # Check if data is available to read
                if sys.stdin in select.select([sys.stdin], [], [], 0)[0]:
                    # Read available data (non-blocking)
                    chunk = sys.stdin.read(chunk_size)
                    if not chunk:  # EOF
                        break
                    buffer += chunk
                    live.update(Markdown(buffer))
                else:
                    # Small delay to prevent busy waiting
                    await asyncio.sleep(0.01)

        except KeyboardInterrupt:
            console.print("\n[yellow]Interrupted by user[/yellow]")
        except Exception as e:
            console.print(f"\n[red]Error: {e}[/red]")
        finally:
            # Ensure final buffer is displayed
            if buffer:
                console.print(Markdown(buffer))


def prettier_code_blocks():
    """Make rich code blocks prettier and easier to copy.

    From https://github.com/samuelcolvin/aicli/blob/v0.8.0/samuelcolvin_aicli.py#L22
    """

    class SimpleCodeBlock(CodeBlock):
        def __rich_console__(
            self, console: Console, options: ConsoleOptions
        ) -> RenderResult:
            code = str(self.text).rstrip()
            yield Text(self.lexer_name, style='dim')
            yield Syntax(
                code,
                self.lexer_name,
                theme=self.theme,
                background_color='default',
                word_wrap=True,
            )
            yield Text(f'/{self.lexer_name}', style='dim')

    Markdown.elements['fence'] = SimpleCodeBlock


if __name__ == '__main__':
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass

# vim: set ft=py: