Update stream_agent.py

stream_agent.py

@@ -2,75 +2,78 @@ from __future__ import annotations
 
 import argparse
 import os
-import time
-from collections.abc import Iterable
+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
 
 
-class StreamingAIAgent:
-    """Small AI-agent with a streaming interface and an offline fallback."""
+@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 __init__(self, model: str | None = None) -> None:
-        self.model = model or os.getenv("OPENAI_MODEL", "openai/gpt-oss-20b:free")
 
-    def invoke(self, question: str) -> str:
-        return "".join(self.stream(question))
-
-    def stream(self, question: str) -> Iterable[str]:
-        client = self._build_client()
-        if client is None:
-            yield from self._offline_stream(question)
-            return
-
-        response = client.chat.completions.create(
-            model=self.model,
-            messages=[
-                {
-                    "role": "system",
-                    "content": "You are a concise educational assistant.",
-                },
-                {"role": "user", "content": question},
-            ],
-            stream=True,
-        )
-        for chunk in response:
-            delta = chunk.choices[0].delta.content
-            if delta:
-                yield delta
-
-    def _build_client(self):
-        api_key = os.getenv("OPENAI_API_KEY")
-        if not api_key:
-            return None
-        try:
-            from openai import OpenAI
-        except ImportError:
-            return None
-        return OpenAI(
-            api_key=api_key,
+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 _offline_stream(self, question: str) -> Iterable[str]:
-        answer = (
-            "Stream mode returns an answer piece by piece instead of waiting "
-            "for the whole response. This makes an AI-agent feel faster and "
-            f"lets the UI show progress while processing: {question}"
+
+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"],
     )
-        for token in answer.split(" "):
-            yield token + " "
-            time.sleep(0.02)
-
+    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 "Explain streaming mode"
-    agent = StreamingAIAgent()
-    for token in agent.stream(question):
-        print(token, end="", flush=True)
-    print()
+    question = " ".join(args.question).strip() or "Сколько стоит молоко и хлеб в Казани?"
+    run_stream(question)
 
 
 if __name__ == "__main__":