diff --git a/main.py b/main.py new file mode 100644 index 0000000..07e4f2f --- /dev/null +++ b/main.py @@ -0,0 +1,115 @@ +"""Stream-режим AI-агента: потоковый вывод токенов и событий updates.""" +from __future__ import annotations + +import os + +from dotenv import load_dotenv +from langchain.agents import create_agent +from langchain.tools import tool +from langchain_openai import ChatOpenAI +from pydantic import SecretStr + +load_dotenv() + +llm = ChatOpenAI( + model=os.getenv("OPENAI_MODEL", "openai/gpt-oss-20b:free"), + base_url=os.getenv("OPENAI_BASE_URL", "https://openrouter.ai/api/v1"), + api_key=SecretStr(os.getenv("OPENAI_API_KEY", "fake")), + temperature=0.7, +) + +step = 1 + + +@tool +def get_price(product: str, city: str) -> str: + """Узнать примерную цену продукта в указанном городе. Возвращает строку таблицы.""" + price_agent = create_agent( + model=llm, + system_prompt=( + "Ты эксперт по розничным ценам в России. " + "Ответ — одна строка таблицы: | Продукт | Цена (руб.) | Магазин |" + ), + ) + result = price_agent.invoke( + { + "messages": [ + { + "role": "human", + "content": ( + f"Какая примерная цена на «{product}» в городе {city}? " + "Верни строку | Продукт | Цена (руб.) | Магазин |" + ), + } + ] + } + ) + return result["messages"][-1].content + + +shopping_agent = create_agent( + model=llm, + tools=[get_price], + system_prompt="Ты помощник по планированию покупок", +) + + +def format_message(message) -> str: + """Текст сообщения или вызов инструмента.""" + if message.content: + return str(message.content) + tool_calls = getattr(message, "tool_calls", None) or [] + if tool_calls: + tc = tool_calls[0] + name = tc.get("name") if isinstance(tc, dict) else getattr(tc, "name", "?") + args = tc.get("args") if isinstance(tc, dict) else getattr(tc, "args", {}) + return f"{name}({args})" + return str(message) + + +def format_chunk_message(chunk_data: tuple) -> None: + """Потоковый вывод токенов с разделителем при смене шага.""" + global step + message, meta = chunk_data + current_step = meta.get("langgraph_step", step) + + if current_step != step: + step = current_step + print("\n --- --- --- \n") + + if message.content: + print(message.content, end="", flush=True) + + +def main() -> None: + global step + step = 1 + + question = ( + "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани." + ) + + stream = shopping_agent.stream( + {"messages": [{"role": "human", "content": question}]}, + 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) + + print() + + +if __name__ == "__main__": + main()