import sys
from typing import Iterable
from langchain_core.messages import HumanMessage, AIMessage
from langchain_openai import ChatOpenAI
from rich.console import Console

console = Console()

def stream_response(messages: list[HumanMessage | AIMessage]) -> Iterable[str]:
    """Yield chunks of the LLM response using streaming."""
    llm = ChatOpenAI(streaming=True, model="gpt-4o-mini")
    chat_history = [msg.model_dump() for msg in messages]
    stream = llm.invoke(chat_history)
    # The stream yields dicts with 'content' key
    for chunk in stream:
        # Each chunk is a dict with 'content'
        content = chunk["content"] if isinstance(chunk, dict) else None
        if content:
            yield content
            console.print(content, end="")

if __name__ == "__main__":
    user_input = sys.argv[1] if len(sys.argv) > 1 else "Hello"
    messages: list[HumanMessage | AIMessage] = [HumanMessage(content=user_input)]
    for _ in stream_response(messages):
        pass