diff --git a/stream_agent.py b/stream_agent.py deleted file mode 100644 index b94737c..0000000 --- a/stream_agent.py +++ /dev/null @@ -1,80 +0,0 @@ -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()