97 lines
2.3 KiB
Python
97 lines
2.3 KiB
Python
"""Console entry point that prints LangGraph agent output in stream mode."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from typing import Any
|
|
|
|
from agent import agent
|
|
|
|
|
|
_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":
|
|
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()
|
|
|
|
|
|
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()
|