From a1fbe22d05ae94595987be3123792dd4f939a3ef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=AD=D0=BC=D0=B8=D0=BB=D1=8C=20=D0=90=D0=BC=D0=B8=D1=80?= =?UTF-8?q?=D0=BE=D0=B2?= Date: Mon, 25 May 2026 21:48:30 +0000 Subject: [PATCH] add main.py --- main.py | 180 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 180 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..070c080 --- /dev/null +++ b/main.py @@ -0,0 +1,180 @@ +"""Агент с памятью разговора и подтверждением вызовов инструментов.""" +from __future__ import annotations + +import os + +from langchain.agents import create_agent +from langchain.tools import tool +from langchain_openai import ChatOpenAI +from langgraph.checkpoint.memory import MemorySaver +from pydantic import SecretStr +from rich.console import Console + +console = Console() + +LM_STUDIO_BASE_URL = os.getenv("LM_STUDIO_BASE_URL", "http://localhost:1234/v1") +LM_STUDIO_MODEL = os.getenv("LM_STUDIO_MODEL", "local-model") +STEP_SEPARATOR = "\n --- --- --- \n" +THREAD_ID = "разговор-1" + + +def build_llm() -> ChatOpenAI: + return ChatOpenAI( + model=LM_STUDIO_MODEL, + base_url=LM_STUDIO_BASE_URL, + api_key=SecretStr(os.getenv("OPENAI_API_KEY", "fake")), + temperature=0.7, + ) + + +def _extract_table(text: str) -> str: + lines = [line for line in text.splitlines() if "|" in line] + return "\n".join(lines) if lines else text.strip() + + +def _build_price_subagent(llm: ChatOpenAI): + return create_agent( + model=llm, + system_prompt=( + "Ты аналитик цен. По продукту, городу и при необходимости дате " + "оцени реалистичную цену в рублях. " + "Ответь одной строкой markdown-таблицы: | Продукт | Цена (руб.) | Магазин |" + ), + ) + + +def make_get_price_tool(llm: ChatOpenAI): + price_subagent = _build_price_subagent(llm) + + @tool + def get_price(product: str, city: str, date: str = "") -> str: + """Возвращает примерную цену продукта в городе (опционально с учётом даты/периода). + + Args: + product: название продукта или категория (молоко, хлеб, погода как запрос цен и т.д.) + city: город + date: период (сегодня, завтра, через неделю) — необязательно + """ + when = f", период: {date}" if date else "" + prompt = ( + f"Город: {city}. Продукт/запрос: {product}{when}. " + "Верни одну строку таблицы с ценой и магазином." + ) + result = price_subagent.invoke( + {"messages": [{"role": "human", "content": prompt}]} + ) + last = result["messages"][-1] + return _extract_table(getattr(last, "content", str(last))) + + return get_price + + +def format_message(message) -> str: + content = getattr(message, "content", None) + if content: + return str(content) + tool_calls = getattr(message, "tool_calls", None) or [] + if tool_calls: + call = tool_calls[0] + name = call.get("name") if isinstance(call, dict) else getattr(call, "name", "") + args = call.get("args") if isinstance(call, dict) else getattr(call, "args", {}) + return f"{name}({args})" + return str(message) + + +def _tool_call_repr(state) -> str: + messages = state.values.get("messages", []) + if not messages: + return "unknown_tool()" + last = messages[-1] + tool_calls = getattr(last, "tool_calls", None) or [] + if not tool_calls: + return "unknown_tool()" + call = tool_calls[0] + name = call.get("name") if isinstance(call, dict) else getattr(call, "name", "") + args = call.get("args") if isinstance(call, dict) else getattr(call, "args", {}) + return f'{name}({args})' + + +def build_agent(llm: ChatOpenAI): + memory = MemorySaver() + get_price = make_get_price_tool(llm) + return create_agent( + model=llm, + tools=[get_price], + system_prompt="Ты помощник по планированию покупок", + checkpointer=memory, + interrupt_before=["tools"], + ) + + +def run_chat() -> None: + llm = build_llm() + agent = build_agent(llm) + config = {"configurable": {"thread_id": THREAD_ID}} + step_holder = {"step": 1} + + def format_chunk_message(chunk) -> None: + message, meta = chunk + graph_step = meta.get("langgraph_step", step_holder["step"]) + if graph_step != step_holder["step"]: + step_holder["step"] = graph_step + console.print(STEP_SEPARATOR, end="") + if message.content: + console.print(message.content, end="") + + def ask_and_run(user_input, cfg) -> None: + stream = agent.stream( + user_input, + config=cfg, + stream_mode=["messages", "updates"], + ) + + for chunk in stream: + state = agent.get_state(cfg) + 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.strip(): + console.print(formatted) + + if "__interrupt__" in chunk_data and state.next == ("tools",): + tool_repr = _tool_call_repr(state) + console.print(STEP_SEPARATOR, end="") + console.print(tool_repr) + console.print( + f"Агент хочет вызвать утилиту {tool_repr}" + ) + answer = input("Разрешить? (Y/n): ") + if answer.lower().strip() == "y": + ask_and_run(None, cfg) + else: + console.print("Отменено") + return + + console.print() + + console.print( + "Чат с агентом (память + подтверждение инструментов). " + "Введите exit для выхода." + ) + + while True: + user_input = input("\nВы: ") + if user_input.strip().lower() == "exit": + break + ask_and_run( + {"messages": [{"role": "human", "content": user_input}]}, + config, + ) + + +if __name__ == "__main__": + run_chat()