diff --git a/samuelcolvin_aicli.py b/samuelcolvin_aicli.py index 9c279b5..1b80853 100644 --- a/samuelcolvin_aicli.py +++ b/samuelcolvin_aicli.py @@ -61,13 +61,16 @@ def cli() -> int: if stream: content = '' markdown = Markdown(content) - with Live(markdown, refresh_per_second=15, console=console) as live: - for chunk in response: - if chunk['choices'][0]['finish_reason'] is not None: - break - chunk_text = chunk['choices'][0]['delta'].get('content', '') - content += chunk_text - live.update(Markdown(content)) + try: + with Live(markdown, refresh_per_second=15, console=console) as live: + for chunk in response: + if chunk['choices'][0]['finish_reason'] is not None: + break + chunk_text = chunk['choices'][0]['delta'].get('content', '') + content += chunk_text + live.update(Markdown(content)) + except KeyboardInterrupt: + console.print('[dim]Interrupted[/dim]') else: content = response['choices'][0]['message']['content'] md = Markdown(content)