diff --git a/main.py b/main.py index 07e4f2f..4301be3 100644 --- a/main.py +++ b/main.py @@ -7,93 +7,141 @@ 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, +BROJS_INFERENCE_URL = "https://platform.brojs.ru/jrnl-bh/api/inference/v1" +DEFAULT_MODEL = "openai/gpt-oss-20b:free" +STEP_SEPARATOR = "\n --- --- --- \n" + +DEFAULT_QUESTION = ( + "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани." ) -step = 1 + +def _api_key() -> str: + return ( + os.getenv("OPENAI_API_KEY") + or os.getenv("JOURNAL_MCP_PAT") + or os.getenv("JOURNAL_TOKEN") + or "" + ) -@tool -def get_price(product: str, city: str) -> str: - """Узнать примерную цену продукта в указанном городе. Возвращает строку таблицы.""" - price_agent = create_agent( +def _base_url() -> str: + if os.getenv("OPENAI_BASE_URL"): + return os.environ["OPENAI_BASE_URL"] + if os.getenv("OPENAI_API_KEY"): + return os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1") + return BROJS_INFERENCE_URL + + +def _model() -> str: + return os.getenv("OPENAI_MODEL") or os.getenv("OPENROUTER_MODEL") or DEFAULT_MODEL + + +def build_llm() -> ChatOpenAI: + return ChatOpenAI( + model=_model(), + base_url=_base_url(), + api_key=_api_key(), + 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" + "| Продукт | Цена (руб.) | Магазин |" ), ) - result = price_agent.invoke( - { - "messages": [ - { - "role": "human", - "content": ( - f"Какая примерная цена на «{product}» в городе {city}? " - "Верни строку | Продукт | Цена (руб.) | Магазин |" - ), - } - ] - } + + +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 build_agent(llm: ChatOpenAI | None = None): + """Главный агент покупок — точка входа для автопроверки.""" + llm = llm or build_llm() + get_price = make_get_price_tool(llm) + return create_agent( + model=llm, + tools=[get_price], + system_prompt="Ты помощник по планированию покупок", ) - 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) + content = getattr(message, "content", None) + if content: + return str(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", {}) + 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 format_chunk_message(chunk_data: tuple) -> None: - """Потоковый вывод токенов с разделителем при смене шага.""" - global step - message, meta = chunk_data - current_step = meta.get("langgraph_step", step) +def run_shopping_assistant_stream( + agent=None, + question: str = DEFAULT_QUESTION, +) -> None: + """Потоковый запуск агента (messages + updates).""" + agent = agent or build_agent() - 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( + stream = agent.stream( {"messages": [{"role": "human", "content": question}]}, stream_mode=["messages", "updates"], ) + step = 1 + + def format_chunk_message(chunk_data: tuple) -> None: + nonlocal step + message, meta = chunk_data + 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 @@ -105,11 +153,15 @@ def main() -> None: if model_update: last_message = model_update["messages"][-1] formatted = format_message(last_message) - if formatted: + if formatted.strip(): print(formatted) print() +def main() -> None: + run_shopping_assistant_stream() + + if __name__ == "__main__": main()