Skip to content

Commit f5283cd

Browse files
authored
Add progress spinner and output tokens as they are generated (#52)
1 parent 29624ca commit f5283cd

File tree

2 files changed

+9
-2
lines changed

2 files changed

+9
-2
lines changed

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ dependencies = [
99
"httpx",
1010
"openai",
1111
"prompt-toolkit",
12+
"yaspin",
1213
"tomlkit",
1314
"tomli >= 1.1.0; python_version < '3.11'"
1415
]

src/shelloracle/shelloracle.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
import sys
66
from pathlib import Path
77

8+
from yaspin import yaspin
9+
from yaspin.spinners import Spinners
810
from prompt_toolkit import PromptSession
911
from prompt_toolkit.application import create_app_session_from_tty
1012
from prompt_toolkit.history import FileHistory
@@ -57,8 +59,12 @@ async def shelloracle() -> None:
5759
default_prompt = os.environ.get("SHOR_DEFAULT_PROMPT")
5860
prompt = await prompt_user(default_prompt)
5961

60-
async for token in provider.generate(prompt):
61-
sys.stdout.write(token)
62+
shell_command = ""
63+
with create_app_session_from_tty(), patch_stdout(raw=True), yaspin() as sp:
64+
async for token in provider.generate(prompt):
65+
shell_command += token
66+
sp.text = shell_command
67+
sys.stdout.write(shell_command)
6268

6369

6470
def cli() -> None:

0 commit comments

Comments
 (0)