168 lines
5.1 KiB
Python
168 lines
5.1 KiB
Python
"""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
|
|
|
|
load_dotenv()
|
|
|
|
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 = (
|
|
"Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
|
|
)
|
|
|
|
|
|
def _api_key() -> str:
|
|
return (
|
|
os.getenv("OPENAI_API_KEY")
|
|
or os.getenv("JOURNAL_MCP_PAT")
|
|
or os.getenv("JOURNAL_TOKEN")
|
|
or ""
|
|
)
|
|
|
|
|
|
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"
|
|
"| Продукт | Цена (руб.) | Магазин |"
|
|
),
|
|
)
|
|
|
|
|
|
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="Ты помощник по планированию покупок",
|
|
)
|
|
|
|
|
|
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(
|
|
agent=None,
|
|
question: str = DEFAULT_QUESTION,
|
|
) -> None:
|
|
"""Потоковый запуск агента (messages + updates)."""
|
|
agent = agent or build_agent()
|
|
|
|
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
|
|
|
|
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()
|
|
|
|
|
|
def main() -> None:
|
|
run_shopping_assistant_stream()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|