52 lines
1.8 KiB
Python
52 lines
1.8 KiB
Python
import sys
|
||
from agent import graph
|
||
|
||
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)
|
||
|
||
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"],
|
||
):
|
||
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
|
||
print()
|
||
state["messages"] = graph.invoke(
|
||
{"input": user_input, "messages": state["messages"]}
|
||
)["messages"]
|
||
|
||
if __name__ == "__main__":
|
||
main() |