From 988ae42444d8aad1865f4b93d9dd6b559b4a3375 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=A0=D0=B8=D0=BD=D0=B0=D1=80=20=D0=9C=D0=B8=D1=80=D0=B7?=
 =?UTF-8?q?=D0=B0=D0=B3=D0=B8=D1=82=D0=BE=D0=B2?= <mirzagitov03@mail.ru>
Date: Thu, 28 May 2026 10:37:50 +0000
Subject: [PATCH] Update stream_agent.py

---
 stream_agent.py | 96 +++++++++++--------------------------------------
 1 file changed, 21 insertions(+), 75 deletions(-)

diff --git a/stream_agent.py b/stream_agent.py
index b94737c..303b471 100644
--- a/stream_agent.py
+++ b/stream_agent.py
@@ -1,80 +1,26 @@
-from __future__ import annotations
-
-import argparse
-import os
-from typing import Any
-
-from langchain_core.messages import AIMessage
-from langchain_core.tools import tool
+import sys
+from typing import Iterable
+from langchain_core.messages import HumanMessage, AIMessage
 from langchain_openai import ChatOpenAI
-from langgraph.prebuilt import create_react_agent
+from rich.console import Console
 
+console = Console()
 
-@tool
-def get_demo_price(product: str, city: str = "Казань") -> str:
-    """Return a deterministic demo price for the requested product."""
-    prices = {
-        "молоко": "89 рублей",
-        "хлеб": "54 рубля",
-        "сыр": "219 рублей",
-    }
-    return f"{product} в городе {city}: {prices.get(product.lower(), 'цена не найдена')}"
-
-
-def build_agent():
-    llm = ChatOpenAI(
-        model=os.getenv("OPENAI_MODEL", "openai/gpt-oss-20b:free"),
-        base_url=os.getenv("OPENAI_BASE_URL") or None,
-        api_key=os.getenv("OPENAI_API_KEY", "not-needed"),
-        temperature=0,
-        streaming=True,
-    )
-    return create_react_agent(llm, tools=[get_demo_price])
-
-
-def format_message(message: AIMessage) -> str:
-    if message.content:
-        return str(message.content)
-    if message.tool_calls:
-        call = message.tool_calls[0]
-        return f"{call['name']}({call['args']})"
-    return ""
-
-
-def format_chunk_message(chunk: tuple[Any, dict[str, Any]], state: dict[str, int]) -> None:
-    message, meta = chunk
-    current_step = int(meta.get("langgraph_step", state["step"]))
-    if current_step != state["step"]:
-        state["step"] = current_step
-        print("\n --- --- --- \n")
-    if getattr(message, "content", None):
-        print(message.content, end="", flush=True)
-
-
-def run_stream(question: str) -> None:
-    agent = build_agent()
-    stream = agent.stream(
-        {"messages": [{"role": "human", "content": question}]},
-        stream_mode=["messages", "updates"],
-    )
-    state = {"step": 1}
-    for chunk_type, chunk_data in stream:
-        if chunk_type == "messages":
-            format_chunk_message(chunk_data, state)
-        elif chunk_type == "updates" and chunk_data.get("model"):
-            last_message = chunk_data["model"]["messages"][-1]
-            formatted = format_message(last_message)
-            if formatted:
-                print(f"\n{formatted}")
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description="Streaming AI-agent demo")
-    parser.add_argument("question", nargs="*", help="User question")
-    args = parser.parse_args()
-
-    question = " ".join(args.question).strip() or "Сколько стоит молоко и хлеб в Казани?"
-    run_stream(question)
-
+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__":
-    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