Add stream_agent.py

stream_agent.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import argparse
+import os
+from typing import Any
+
+from langchain_core.messages import AIMessage
+from langchain_core.tools import tool
+from langchain_openai import ChatOpenAI
+from langgraph.prebuilt import create_react_agent
+
+
+@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)
+
+
+if __name__ == "__main__":
+    main()