import asyncio
from types import AsyncGeneratorType
from dotenv import load_dotenv
import os
import openai

load_dotenv()

openai.api_key = os.getenv("OPENAI_API_KEY")

async def test_completion_async_stream():
    completion = await openai.Completion.acreate(
        model="text-davinci-003",
        prompt="Hello!",
        stream=True,
    )

    # Assert response body
    assert isinstance(completion, AsyncGeneratorType)

    collected_chunks = []
    completion_text = ''
    # iterate through the stream, if it breaks, the test failed
    async for chunk in completion:
        collected_chunks.append(chunk)
        finish_reason = chunk['choices'][0]['finish_reason'] # type: ignore
        if finish_reason is not None:
            break
        chunk_text = chunk['choices'][0]['text'] # type: ignore
        completion_text += chunk_text  # append the text
    print(completion_text)

if __name__ == "__main__":
    asyncio.run(test_completion_async_stream())