diff --git a/main.py b/main.py new file mode 100644 index 0000000..c91dabe --- /dev/null +++ b/main.py @@ -0,0 +1,65 @@ +import os +from langchain_openai import ChatOpenAI +from langchain_core.messages import HumanMessage +from langchain.tools import tool +from langgraph.checkpoint.memory import MemorySaver +from langgraph.types import Command + +# LLM setup +llm = ChatOpenAI( + model="openai/gpt-oss-20b:free", + base_url="https://openrouter.ai/api/v1", + api_key=os.getenv("OPENAI_API_KEY"), + temperature=0.0, +) + +# Dummy tool: get_price +@tool +def get_price(product: str, city: str) -> str: + """Return a markdown table with price for a product in a city.""" + # In a real scenario, fetch from API. Here we return a static example. + return f"| Продукт | Цена (руб.) | Магазин |\n| {product} | 89 | Магнит |" + +# Agent definition +from langchain.agents import create_agent + +agent = create_agent( + llm=llm, + tools=[get_price], + system_prompt="You are a helpful assistant that can call get_price to provide price information.", +) + +# Stream execution +stream = agent.stream( + { + "messages": [HumanMessage(content="Покажи цену молока и хлеба в Казани.")] + }, + stream_mode=["messages", "updates"], +) + +step = 1 + +def format_chunk_message(chunk): + message, meta = chunk + global step + if meta.get("langgraph_step") != step: + step = meta.get("langgraph_step") + print("\n --- --- --- \n") + if message.content: + print(message.content, end="", flush=True) + +def format_message(message): + if message.content: + return message.content + return f"{message.tool_calls[0]['name']}({message.tool_calls[0]['args']})" + +for chunk in stream: + chunk_type, chunk_data = chunk + if chunk_type == "messages": + format_chunk_message(chunk_data) + elif chunk_type == "updates": + if chunk_data.get("model"): + last_message = chunk_data["model"]["messages"][-1] + print(format_message(last_message)) + +print("\n--- Завершено ---")