Публикация решения: update main.py

main.py

@@ -1,96 +1,49 @@
-"""Console entry point that prints LangGraph agent output in stream mode."""
+"""Console demo for stream-mode LangGraph agent."""
 
-from __future__ import annotations
-
-import sys
+import os
 from typing import Any
 
+from langchain_core.messages import HumanMessage, AIMessage
+from langgraph.graph import StateGraph
+
 from agent import agent
 
-
-_current_step: int | None = None
-
+# Helper to format messages
 
 def format_message(message: Any) -> str:
-    """Format a completed model message or a tool-call request."""
-
-    content = getattr(message, "content", "")
-    if content:
-        return str(content)
-
-    tool_calls = getattr(message, "tool_calls", None) or []
-    if tool_calls:
-        call = tool_calls[0]
+    if hasattr(message, "content") and message.content:
+        return message.content
+    # fallback for tool calls
+    if hasattr(message, "tool_calls") and message.tool_calls:
+        call = message.tool_calls[0]
         return f"{call['name']}({call['args']})"
-
     return ""
 
+# Stream the conversation
 
-def format_chunk_message(chunk: tuple[Any, dict[str, Any]]) -> None:
-    """Print token chunks and separate LangGraph steps."""
-
-    global _current_step
-
-    message, meta = chunk
-    langgraph_step = meta.get("langgraph_step")
-    if _current_step is None:
-        _current_step = langgraph_step
-    elif langgraph_step != _current_step:
-        _current_step = langgraph_step
-        print("\n--- --- ---\n", end="")
-
-    if message.content:
-        print(message.content, end="", flush=True)
-
-
-def run_stream(query: str) -> None:
-    """Run the agent with .stream() and handle messages/updates chunks."""
-
-    global _current_step
-    _current_step = None
-
+def main() -> None:
+    # initial human message
+    human_msg = HumanMessage(content="Составь список покупок: молоко, хлеб, яблоки")
+    # start stream
     stream = agent.stream(
-        {"messages": [{"role": "human", "content": query}]},
+        {"messages": [human_msg]},
         stream_mode=["messages", "updates"],
     )
 
+    step = 1
     for chunk in stream:
         chunk_type, chunk_data = chunk
-
         if chunk_type == "messages":
-            format_chunk_message(chunk_data)
-
-        if chunk_type == "updates":
-            model_update = chunk_data.get("model")
-            if model_update:
-                last_message = model_update["messages"][-1]
-                formatted = format_message(last_message)
-                if formatted:
-                    print(formatted, end="", flush=True)
-
-    print()
-
-
-def main() -> None:
-    """Read one CLI query or start a small interactive loop."""
-
-    if len(sys.argv) > 1:
-        run_stream(" ".join(sys.argv[1:]))
-        return
-
-    print("Stream-mode AI agent. Type 'exit' to quit.")
-    while True:
-        try:
-            query = input("You: ").strip()
-        except EOFError:
-            print()
-            break
-
-        if query.lower() in {"exit", "quit"}:
-            break
-        if query:
-            run_stream(query)
-
+            message, meta = chunk_data
+            if meta.get("langgraph_step") != step:
+                step = meta.get("langgraph_step")
+                print("\n--- --- ---\n")
+            if message.content:
+                print(message.content, end="", flush=True)
+        elif chunk_type == "updates":
+            if chunk_data.get("model"):
+                last_msg = chunk_data["model"]["messages"][-1]
+                print(format_message(last_msg))
 
 if __name__ == "__main__":
     main()