diff --git a/main.py b/main.py new file mode 100644 index 0000000..440d1a7 --- /dev/null +++ b/main.py @@ -0,0 +1,133 @@ +"""Иерархический AI-агент со stream-выводом (LangChain + LM Studio).""" +from __future__ import annotations + +import os + +from langchain.agents import create_agent +from langchain.tools import tool +from langchain_openai import ChatOpenAI +from pydantic import SecretStr + +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" + + +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] + if lines: + return "\n".join(lines) + return text.strip() + + +def _build_price_subagent(llm: ChatOpenAI): + return create_agent( + model=llm, + system_prompt=( + "Ты аналитик цен на продукты питания. " + "По названию продукта и городу оцени реалистичную цену в рублях, " + "опираясь на типичные российские розничные цены. " + "Ответь ТОЛЬКО одной строкой markdown-таблицы в формате:\n" + "| Продукт | Цена (руб.) | Магазин |" + ), + ) + + +def make_get_price_tool(llm: ChatOpenAI): + price_subagent = _build_price_subagent(llm) + + @tool + def get_price(product: str, city: str) -> str: + """Возвращает примерную цену продукта в указанном городе. + + Args: + product: название продукта (молоко, хлеб, яблоки и т.д.) + city: город покупателя + """ + prompt = ( + f"Город: {city}. Продукт: {product}. " + "Верни одну строку таблицы с реалистичной ценой и названием магазина." + ) + result = price_subagent.invoke( + {"messages": [{"role": "human", "content": prompt}]} + ) + last = result["messages"][-1] + content = getattr(last, "content", str(last)) + return _extract_table(content) + + 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 run_shopping_assistant_stream() -> None: + llm = build_llm() + get_price = make_get_price_tool(llm) + + main_agent = create_agent( + model=llm, + tools=[get_price], + system_prompt="Ты помощник по планированию покупок", + ) + + question = ( + "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани." + ) + + stream = main_agent.stream( + {"messages": [{"role": "human", "content": question}]}, + stream_mode=["messages", "updates"], + ) + + step = 1 + + def format_chunk_message(chunk) -> None: + nonlocal step + message, meta = chunk + graph_step = meta.get("langgraph_step", step) + if graph_step != step: + step = graph_step + print(STEP_SEPARATOR, end="") + if message.content: + print(message.content, end="", flush=True) + + 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.strip(): + print(formatted) + + print() + + +if __name__ == "__main__": + run_shopping_assistant_stream()