125 lines
5.4 KiB
Python
125 lines
5.4 KiB
Python
import os
|
||
import asyncio
|
||
from rich.console import Console
|
||
from rich.text import Text
|
||
from langchain_openai import ChatOpenAI
|
||
from langchain_core.messages import HumanMessage
|
||
from langchain.tools import tool
|
||
from deepagents import create_deep_agent
|
||
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
|
||
from langgraph.checkpoint.memory import MemorySaver
|
||
|
||
# ---------- Консоль ----------
|
||
console = Console()
|
||
|
||
# ---------- LLM ----------
|
||
llm = ChatOpenAI(
|
||
model="openai/gpt-oss-20b:free",
|
||
base_url="https://openrouter.ai/api/v1",
|
||
api_key=os.getenv("OPENAI_API_KEY"),
|
||
temperature=0.0,
|
||
)
|
||
|
||
# ---------- Инструменты ----------
|
||
@tool
|
||
def get_price(city: str, date: str) -> str:
|
||
"""
|
||
Возвращает условную цену погоды для указанного города и даты.
|
||
Параметры:
|
||
city: название города
|
||
date: дата в произвольном формате (например, "сегодня", "завтра")
|
||
"""
|
||
# Здесь могла бы быть реальная логика обращения к API погоды.
|
||
return f"Цена погоды в {city} на {date}: 42 условных единиц."
|
||
|
||
# ---------- Бэкенд ----------
|
||
backend = CompositeBackend(
|
||
[
|
||
LocalShellBackend(workspace_dir="./workspace"),
|
||
FilesystemBackend(),
|
||
]
|
||
)
|
||
|
||
# ---------- Память ----------
|
||
memory = MemorySaver()
|
||
|
||
# ---------- Агент ----------
|
||
agent = create_deep_agent(
|
||
model=llm,
|
||
tools=[get_price],
|
||
backend=backend,
|
||
system_prompt="You are a helpful assistant. Use the provided tools when needed.",
|
||
checkpointer=memory,
|
||
interrupt_before=["tools"], # пауза перед каждым вызовом инструмента
|
||
)
|
||
|
||
# ---------- Функция обработки одного запроса ----------
|
||
async def ask_and_run(user_input, config):
|
||
"""
|
||
user_input: dict вида {"messages": [{"role": "human", "content": "..."}]}
|
||
или None (для возобновления после подтверждения)
|
||
config: dict с ключом "configurable" -> {"thread_id": "..."}
|
||
"""
|
||
# Если передан None, просто продолжаем без новых сообщений
|
||
input_payload = user_input if user_input is not None else {}
|
||
|
||
async for chunk in agent.astream(
|
||
input_payload,
|
||
config,
|
||
stream_mode=["messages", "updates"],
|
||
):
|
||
# Получаем актуальное состояние после каждого чанка
|
||
state = agent.get_state(config)
|
||
|
||
chunk_type, chunk_data = chunk
|
||
|
||
# 1. Вывод сообщений (поток токенов)
|
||
if chunk_type == "messages":
|
||
# chunk_data содержит список сообщений, но мы выводим только новые токены
|
||
# Для простоты выводим полное содержание последнего сообщения, если оно от ассистента
|
||
last_msg = chunk_data[-1]
|
||
if getattr(last_msg, "role", None) == "assistant":
|
||
console.print(Text(last_msg.content, style="green"))
|
||
|
||
# 2. Вывод обновлений (инструменты)
|
||
if chunk_type == "updates":
|
||
# Здесь можно вывести информацию о вызове инструмента, если она есть
|
||
pass # в данном простом примере ничего не делаем
|
||
|
||
# 3. Обработка прерывания перед инструментом
|
||
if "__interrupt__" in chunk_data and state.next == ("tools",):
|
||
# Получаем информацию о запланированном вызове инструмента
|
||
tool_msg = state.values["messages"][-1]
|
||
tool_call = tool_msg.tool_calls[0]
|
||
tool_name = tool_call["name"]
|
||
tool_args = tool_call["args"]
|
||
console.print(f"[bold yellow]Agent wants to call {tool_name}{tool_args}[/]")
|
||
|
||
answer = input("Разрешить? (Y/n): ").strip().lower()
|
||
if answer == "n" or answer == "no":
|
||
console.print("[red]Отменено[/]")
|
||
# Прерываем текущий поток, не вызывая инструмент
|
||
break
|
||
else:
|
||
# Возобновляем поток без нового пользовательского сообщения
|
||
await ask_and_run(None, config)
|
||
# После возврата продолжаем обработку (обычно поток уже завершён)
|
||
return
|
||
|
||
# ---------- Основной цикл ----------
|
||
async def main():
|
||
thread_id = "conversation-1"
|
||
config = {"configurable": {"thread_id": thread_id}}
|
||
|
||
console.print("[bold cyan]=== Чат с подтверждением инструментов ===[/]")
|
||
while True:
|
||
user_text = input("\nВы: ").strip()
|
||
if user_text.lower() in {"exit", "quit"}:
|
||
console.print("[bold magenta]Завершение работы.[/]")
|
||
break
|
||
|
||
user_message = {"messages": [{"role": "human", "content": user_text}]}
|
||
await ask_and_run(user_message, config)
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main()) |