119 lines
5.4 KiB
Python
119 lines
5.4 KiB
Python
from langchain_openai import ChatOpenAI
|
||
from pydantic import SecretStr
|
||
from langchain.agents import create_agent
|
||
from langchain.tools import tool
|
||
from langgraph.checkpoint.memory import MemorySaver
|
||
from rich.console import Console
|
||
|
||
# Инициализация LLM с использованием плейсхолдеров
|
||
llm = ChatOpenAI(
|
||
model="google/gemma-4-26b-a4b",
|
||
base_url="http://192.168.0.120:1234/v1",
|
||
api_key=SecretStr("lm-studio"),
|
||
temperature=0.7,
|
||
)
|
||
|
||
console = Console()
|
||
|
||
# Определение инструмента
|
||
@tool
|
||
def get_price(city: str, date: str):
|
||
"""Возвращает прогноз погоды (цены/состояние) для указанного города и даты."""
|
||
# Имитация логики
|
||
return f"В городе {city} на дату {date} ожидается солнечная погода, +20°C."
|
||
|
||
tools = [get_price]
|
||
|
||
# Настройка памяти и агента с механизмом прерывания (interrupt)
|
||
memory = MemorySaver()
|
||
agent = create_agent(
|
||
model=llm,
|
||
tools=tools,
|
||
system_prompt="Ты полезный помощкиник. Если пользователь спрашивает о погоде, используй инструмент get_price.",
|
||
checkpointer=memory,
|
||
interrupt_before=['tools'],
|
||
)
|
||
|
||
# Конфигурация потока (thread_id обеспечивает память разговора)
|
||
config = {"configurable": {"thread_id": "chat-session-123"}}
|
||
|
||
def ask_and_run(user_input, config):
|
||
"""Основная функция обработки сообщений и управления циклом подтверждения."""
|
||
# Если user_input is None, мы просто продолжаем выполнение (возобновление после паузы)
|
||
input_payload = {"messages": [{"role": "human", "content": user_input}]} if user_input else None
|
||
|
||
# stream_mode=['messages', 'updates'] позволяет обрабатывать и токены, и изменения состояния
|
||
stream = agent.stream(input_payload, config=config, stream_mode=['messages', 'updates'])
|
||
|
||
try:
|
||
for chunk in stream:
|
||
chunk_type, chunk_data = chunk
|
||
state = agent.get_state(config)
|
||
|
||
if chunk_type == 'messages':
|
||
# Потоковый вывод текста (токены)
|
||
# В новых версиях LangGraph/LangChain чанк сообщения может быть объектом BaseMessage
|
||
content = ""
|
||
if hasattr(chunk_data, 'content'):
|
||
content = chunk_data.content
|
||
elif isinstance(chunk_data, str):
|
||
content = chunk_data
|
||
|
||
if content:
|
||
console.print(content, end="", flush=True)
|
||
|
||
elif chunk_type == 'updates':
|
||
# Здесь можно обрабатывать обновления узлов, если нужно
|
||
pass
|
||
|
||
# Проверка на прерывание (интеррупт перед инструментом)
|
||
# В LangGraph при interrupt_before ключ '__interrupt__' появляется в данных чанка
|
||
if '__interrupt__' in chunk_data and state.next == ('tools',):
|
||
console.print("\n" + "---" * 10)
|
||
|
||
# Извлекаем информацию о том, какой инструмент вызывается
|
||
last_message = state.values['messages'][-1]
|
||
if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
|
||
tool_call = last_message.tool_calls[0]
|
||
console.print(f"{tool_call['name']}({tool_call['args']})")
|
||
console.print(f"Агент хочет вызвать утилиту {tool_call['name']}({tool_call['args']})")
|
||
|
||
answer = input("Разрешить? (Y/n): ")
|
||
|
||
if answer.lower().strip() == 'y':
|
||
# Рекурсивный вызов с None для продолжения выполнения
|
||
ask_and_run(None, config)
|
||
else:
|
||
console.print("Отменено")
|
||
return # Выход из текущей итерации функции
|
||
|
||
# Печать переноса строки после завершения стрима (если не было прерывания)
|
||
if user_input is not None:
|
||
console.print()
|
||
|
||
except Exception as e:
|
||
# Обработка ошибок или остановки потока
|
||
pass
|
||
|
||
def main():
|
||
console.print("[bold blue]Чат запущен. Напишите 'exit' для выхода.[/bold blue]")
|
||
|
||
while True:
|
||
try:
|
||
user_text = input("\nВы: ")
|
||
if user_text.lower().strip() == 'exit':
|
||
console.print("До свидания!")
|
||
break
|
||
|
||
# Запуск агента с новым сообщением
|
||
ask_and_run(
|
||
{"messages": [{"role": "human", "content": user_text}]},
|
||
config,
|
||
)
|
||
except EOFError:
|
||
break
|
||
except Exception as e:
|
||
console.print(f"Ошибка: {e}")
|
||
|
||
if __name__ == "__main__":
|
||
main() |