"""Simple streaming AI agent using OpenAI ChatCompletion API.

The module exposes a single function ``stream_chat`` that yields
text fragments as they arrive from the model.  It is intentionally
light‑weight so it can be used as a drop‑in component in larger
applications.

Dependencies
------------
* ``openai`` – official OpenAI SDK

Environment
-----------
An OpenAI API key must be available either as the environment
variable ``OPENAI_API_KEY`` or passed explicitly via the
``api_key`` argument.
"""

from __future__ import annotations

import os
from typing import Iterable, Generator, Dict, List, Any

import openai

# Ensure the OpenAI key is set when the module is imported.
# Users can override by passing ``api_key`` to ``stream_chat``.
openai.api_key = os.getenv("OPENAI_API_KEY")


def stream_chat(
    messages: List[Dict[str, str]],
    *,
    model: str = "gpt-3.5-turbo",
    temperature: float = 0.7,
    api_key: str | None = None,
) -> Generator[str, None, None]:
    """Yield model output token by token.

    Parameters
    ----------
    messages:
        A list of message objects compatible with the ChatCompletion API.
    model:
        The model to use. Defaults to ``gpt-3.5-turbo``.
    temperature:
        Sampling temperature. Defaults to 0.7.
    api_key:
        Optional API key. If provided it overrides the environment
        variable.

    Yields
    ------
    str
        The next fragment of the assistant's reply.
    """

    if api_key is not None:
        openai.api_key = api_key

    # The SDK returns an iterator over chunks when stream=True.
    response = openai.ChatCompletion.create(
        model=model,
        messages=messages,
        temperature=temperature,
        stream=True,
    )

    # The response is an iterator of chunks. Each chunk contains a
    # ``choices[0].delta`` dict with the text fragment.
    for chunk in response:
        try:
            delta = chunk["choices"][0]["delta"]
            if "content" in delta:
                yield delta["content"]
        except Exception as exc:  # pragma: no cover – defensive
            # In a real application you might log this.
            print(f"Error processing chunk: {exc}")
            continue


# If the module is executed directly, run a small demo.
if __name__ == "__main__":  # pragma: no cover
    demo_messages = [
        {"role": "user", "content": "Write a short poem about the ocean."}
    ]
    print("Streaming response: ")
    for token in stream_chat(demo_messages):
        print(token, end="", flush=True)
    print()