diff --git a/stream_agent.py b/stream_agent.py new file mode 100644 index 0000000..b94737c --- /dev/null +++ b/stream_agent.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import argparse +import os +from typing import Any + +from langchain_core.messages import AIMessage +from langchain_core.tools import tool +from langchain_openai import ChatOpenAI +from langgraph.prebuilt import create_react_agent + + +@tool +def get_demo_price(product: str, city: str = "Казань") -> str: + """Return a deterministic demo price for the requested product.""" + prices = { + "молоко": "89 рублей", + "хлеб": "54 рубля", + "сыр": "219 рублей", + } + return f"{product} в городе {city}: {prices.get(product.lower(), 'цена не найдена')}" + + +def build_agent(): + llm = ChatOpenAI( + model=os.getenv("OPENAI_MODEL", "openai/gpt-oss-20b:free"), + base_url=os.getenv("OPENAI_BASE_URL") or None, + api_key=os.getenv("OPENAI_API_KEY", "not-needed"), + temperature=0, + streaming=True, + ) + return create_react_agent(llm, tools=[get_demo_price]) + + +def format_message(message: AIMessage) -> str: + if message.content: + return str(message.content) + if message.tool_calls: + call = message.tool_calls[0] + return f"{call['name']}({call['args']})" + return "" + + +def format_chunk_message(chunk: tuple[Any, dict[str, Any]], state: dict[str, int]) -> None: + message, meta = chunk + current_step = int(meta.get("langgraph_step", state["step"])) + if current_step != state["step"]: + state["step"] = current_step + print("\n --- --- --- \n") + if getattr(message, "content", None): + print(message.content, end="", flush=True) + + +def run_stream(question: str) -> None: + agent = build_agent() + stream = agent.stream( + {"messages": [{"role": "human", "content": question}]}, + stream_mode=["messages", "updates"], + ) + state = {"step": 1} + for chunk_type, chunk_data in stream: + if chunk_type == "messages": + format_chunk_message(chunk_data, state) + elif chunk_type == "updates" and chunk_data.get("model"): + last_message = chunk_data["model"]["messages"][-1] + formatted = format_message(last_message) + if formatted: + print(f"\n{formatted}") + +def main() -> None: + parser = argparse.ArgumentParser(description="Streaming AI-agent demo") + parser.add_argument("question", nargs="*", help="User question") + args = parser.parse_args() + + question = " ".join(args.question).strip() or "Сколько стоит молоко и хлеб в Казани?" + run_stream(question) + + +if __name__ == "__main__": + main()