Submit solution for task 699cc158d6d3a5544a3ed35b: update main.py

main.py

@@ -1,18 +1,31 @@
 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:
-        user_input = line.rstrip('\n')
+    # Если передан аргумент, используем его как сообщение; иначе читаем из 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
-        # Stream output with separators between steps
+            process(user_input, state)
+
+def process(user_input, state):
     current_step = None
-        for chunk_type, chunk_data in graph.stream({"input": user_input, "messages": state["messages"]}, stream_mode=["messages", "updates"]):
+    for chunk_type, chunk_data in graph.stream(
+        {"input": user_input, "messages": state["messages"]},
+        stream_mode=["messages", "updates"],
+    ):
         if chunk_type == "messages":
             message, meta = chunk_data
             step = meta.get("langgraph_step")
@@ -25,12 +38,15 @@ if __name__ == "__main__":
                 print(message.content, end="")
         elif chunk_type == "updates":
             last_message = chunk_data.get("model", {}).get("messages", [-1])[-1]
-                # If the last message is an AIMessage with content, print it
             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"]
-    print("Goodbye!")
+    state["messages"] = graph.invoke(
+        {"input": user_input, "messages": state["messages"]}
+    )["messages"]
+
+if __name__ == "__main__":
+    main()