import os

from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.dataclasses import ChatMessage, StreamingChunk
from haystack.utils import Secret

received = []


def show_chunk(chunk: StreamingChunk) -> None:
    if chunk.content:
        received.append(chunk.content)
        print(chunk.content, end="", flush=True)


generator = OpenAIChatGenerator(
    api_key=Secret.from_env_var("OPENAI_API_KEY"),
    api_base_url=os.getenv("OPENAI_API_BASE_URL"),
    model=os.getenv("OPENAI_MODEL", "gpt-5-mini"),
    streaming_callback=show_chunk,
    generation_kwargs={"n": 1},
)

print("stream: ", end="")
result = generator.run(
    messages=[
        ChatMessage.from_user(
            "Explain Haystack response streaming in one short sentence."
        )
    ]
)
reply = result["replies"][0]
streamed_text = "".join(received)

print(f"\nfinal: {reply.text}")
print(f"match: {streamed_text == reply.text}")