diff --git a/main.py b/main.py
index a0b8429..8a2225c 100644
--- a/main.py
+++ b/main.py
@@ -1,52 +1,96 @@
+"""Console entry point that prints LangGraph agent output in stream mode."""
+
+from __future__ import annotations
+
 import sys
-from agent import graph
+from typing import Any
 
-def main():
-    print("LangGraph stream‑mode demo. Type 'exit' to quit.")
-    state = {"messages": []}
-    # Если передан аргумент, используем его как сообщение; иначе читаем из stdin
-    if len(sys.argv) > 1:
-        user_input = " ".join(sys.argv[1:])
-        process(user_input, state)
-    else:
-        while True:
-            try:
-                user_input = input("You: ")
-            except EOFError:
-                break
-            if not user_input:
-                continue
-            if user_input.lower() in {"exit", "quit"}:
-                break
-            process(user_input, state)
+from agent import agent
 
-def process(user_input, state):
-    current_step = None
-    for chunk_type, chunk_data in graph.stream(
-        {"input": user_input, "messages": state["messages"]},
+
+_current_step: int | None = None
+
+
+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]
+        return f"{call['name']}({call['args']})"
+
+    return ""
+
+
+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
+
+    stream = agent.stream(
+        {"messages": [{"role": "human", "content": query}]},
         stream_mode=["messages", "updates"],
-    ):
+    )
+
+    for chunk in stream:
+        chunk_type, chunk_data = chunk
+
         if chunk_type == "messages":
-            message, meta = chunk_data
-            step = meta.get("langgraph_step")
-            if current_step is None:
-                current_step = step
-            elif step != current_step:
-                print("\n--- --- ---\n", end="")
-                current_step = step
-            if message.content:
-                print(message.content, end="")
-        elif chunk_type == "updates":
-            last_message = chunk_data.get("model", {}).get("messages", [-1])[-1]
-            try:
-                if hasattr(last_message, "content") and last_message.content:
-                    print(last_message.content, end="")
-            except Exception:
-                pass
+            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()
-    state["messages"] = graph.invoke(
-        {"input": user_input, "messages": state["messages"]}
-    )["messages"]
+
+
+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)
+
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()