Skip to content

Commit

Permalink
Add progress spinner and output tokens as they are generated (#52)
Browse files Browse the repository at this point in the history
  • Loading branch information
djcopley authored Feb 13, 2024
1 parent 29624ca commit f5283cd
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 2 deletions.
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ dependencies = [
"httpx",
"openai",
"prompt-toolkit",
"yaspin",
"tomlkit",
"tomli >= 1.1.0; python_version < '3.11'"
]
Expand Down
10 changes: 8 additions & 2 deletions src/shelloracle/shelloracle.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
import sys
from pathlib import Path

from yaspin import yaspin
from yaspin.spinners import Spinners
from prompt_toolkit import PromptSession
from prompt_toolkit.application import create_app_session_from_tty
from prompt_toolkit.history import FileHistory
Expand Down Expand Up @@ -57,8 +59,12 @@ async def shelloracle() -> None:
default_prompt = os.environ.get("SHOR_DEFAULT_PROMPT")
prompt = await prompt_user(default_prompt)

async for token in provider.generate(prompt):
sys.stdout.write(token)
shell_command = ""
with create_app_session_from_tty(), patch_stdout(raw=True), yaspin() as sp:
async for token in provider.generate(prompt):
shell_command += token
sp.text = shell_command
sys.stdout.write(shell_command)


def cli() -> None:
Expand Down

0 comments on commit f5283cd

Please sign in to comment.