Issues Handling ChatGPT Streaming Response in Terminal using OpenAI API - Using Python, rich library

Question

I am trying to integrate the openAi API model - gpt-4 with Terminal to enable ChatGPT. My objective is to receive streaming responses from ChatGPT and print them in the Terminal. Although I can successfully print the entire response without streaming, I'm facing issues with streaming responses. Specifically, the ask_stream function is printing every word on a new line, which is not the desired behavior. I'm using the rich library to handle Markups

My code:

import openai
from rich.markdown import Markdown
from rich.console import Console
from prompt_toolkit import PromptSession
from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
from prompt_toolkit.completion import WordCompleter
from prompt_toolkit.history import InMemoryHistory
import time
import argparse
import asyncio

openai.api_key = "MY API KEY"
model = "gpt-4"
delay_time = 0.01
max_response_length = 200
console = Console()


async def ask_stream(prompt):
    response = openai.ChatCompletion.create(model='gpt-4',
                                            messages=[{"role": "user", "content": f"{prompt}"}], max_tokens=8000,
                                            temperature=0.4, stream=True)
    answer = ''
    for event in response:
        if answer:
            console.print(Markdown(answer), end='')
        # sys.stdout.flush()
        event_text = event['choices'][0]['delta']
        answer = event_text.get('content', '')
        time.sleep(0.01)


async def ask(prompt) -> Markdown:
    if prompt:
        completion = openai.ChatCompletion.create(model=model,
                                                  messages=[{"role": "user", "content": f"{prompt}"}])
        if completion:
            if 'error' in completion:
                return completion['error']['message']
            return Markdown(completion.choices[0].message.content)
        else:
            raise Exception("")


def create_session() -> PromptSession:
    return PromptSession(history=InMemoryHistory())


async def get_input_async(
        session: PromptSession = None,
        completer: WordCompleter = None,
) -> str:
    """
    Multiline input function.
    """
    return await session.prompt_async(
        completer=completer,
        multiline=True,
        auto_suggest=AutoSuggestFromHistory(),
    )


async def main():
    print(f"Starting Chatgpt with model - {model}")
    session = create_session()
    while True:
        print("\nYou:")
        question = await get_input_async(session=session)
        print()
        print()
        if question == "!exit":
            break
        elif question == "!help":
            print(
                """
            !help - Show this help message
            !exit - Exit the program
            """,
            )
            continue
        print("ChatGPT:")
        if args.no_stream:
            console.print(await ask(prompt=question))
        else:
            await ask_stream(prompt=question)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--no-stream", action="store_true")
    args = parser.parse_args()
    asyncio.run(main())

ask_stream prints like below

Can someone suggest a solution to fix this issue? I am pretty new to Python.

Answer 1

You should change this part of your code:

    answer = ''
    for event in response:
        if answer:
            console.print(Markdown(answer), end='')
        # sys.stdout.flush()
        event_text = event['choices'][0]['delta']
        answer = event_text.get('content', '')
        time.sleep(0.01)

To this one:

    answer = []
    for event in response:
        # sys.stdout.flush()
        event_text = event['choices'][0]['delta']
        answer.append(event_text.get('content', ''))

        if answer:
            console.clear()
            console.print(Markdown(' '.join(answer)), end='')

        time.sleep(0.01)