50 lines
1.5 KiB
Python
50 lines
1.5 KiB
Python
"""Console demo for stream-mode LangGraph agent."""
|
|
|
|
import os
|
|
from typing import Any
|
|
|
|
from langchain_core.messages import HumanMessage, AIMessage
|
|
from langgraph.graph import StateGraph
|
|
|
|
from agent import agent
|
|
|
|
# Helper to format messages
|
|
|
|
def format_message(message: Any) -> str:
|
|
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 main() -> None:
|
|
# initial human message
|
|
human_msg = HumanMessage(content="Составь список покупок: молоко, хлеб, яблоки")
|
|
# start stream
|
|
stream = agent.stream(
|
|
{"messages": [human_msg]},
|
|
stream_mode=["messages", "updates"],
|
|
)
|
|
|
|
step = 1
|
|
for chunk in stream:
|
|
chunk_type, chunk_data = chunk
|
|
if chunk_type == "messages":
|
|
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()
|