Submit solution for task 699cc158d6d3a5544a3ed35b: update main.py

main.py

@@ -1,36 +1,52 @@
 import sys
 from agent import graph
 
-if __name__ == "__main__":
+def main():
     print("LangGraph stream‑mode demo. Type 'exit' to quit.")
     state = {"messages": []}
-    for line in sys.stdin:
+    # Если передан аргумент, используем его как сообщение; иначе читаем из stdin
-        user_input = line.rstrip('\n')
+    if len(sys.argv) > 1:
-        if not user_input:
+        user_input = " ".join(sys.argv[1:])
-            continue
+        process(user_input, state)
-        if user_input.lower() in {"exit", "quit"}:
+    else:
-            break
+        while True:
-        # Stream output with separators between steps
+            try:
-        current_step = None
+                user_input = input("You: ")
-        for chunk_type, chunk_data in graph.stream({"input": user_input, "messages": state["messages"]}, stream_mode=["messages", "updates"]):
+            except EOFError:
-            if chunk_type == "messages":
+                break
-                message, meta = chunk_data
+            if not user_input:
-                step = meta.get("langgraph_step")
+                continue
-                if current_step is None:
+            if user_input.lower() in {"exit", "quit"}:
-                    current_step = step
+                break
-                elif step != current_step:
+            process(user_input, state)
-                    print("\n--- --- ---\n", end="")
+
-                    current_step = step
+def process(user_input, state):
-                if message.content:
+    current_step = None
-                    print(message.content, end="")
+    for chunk_type, chunk_data in graph.stream(
-            elif chunk_type == "updates":
+        {"input": user_input, "messages": state["messages"]},
-                last_message = chunk_data.get("model", {}).get("messages", [-1])[-1]
+        stream_mode=["messages", "updates"],
-                # If the last message is an AIMessage with content, print it
+    ):
-                try:
+        if chunk_type == "messages":
-                    if hasattr(last_message, "content") and last_message.content:
+            message, meta = chunk_data
-                        print(last_message.content, end="")
+            step = meta.get("langgraph_step")
-                except Exception:
+            if current_step is None:
-                    pass
+                current_step = step
-        print()
+            elif step != current_step:
-        state["messages"] = graph.invoke({"input": user_input, "messages": state["messages"]})["messages"]
+                print("\n--- --- ---\n", end="")
-    print("Goodbye!")
+                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
+    print()
+    state["messages"] = graph.invoke(
+        {"input": user_input, "messages": state["messages"]}
+    )["messages"]
+
+if __name__ == "__main__":
+    main()