27 lines
995 B
Python
27 lines
995 B
Python
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
|