add main.py
This commit is contained in:
@@ -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()
|
||||
Reference in New Issue
Block a user