470 lines
19 KiB
Python
470 lines
19 KiB
Python
"""
|
||
Этот файл содержит узлы для графа состояний агента.
|
||
Каждый узел выполняет определенную функцию, такую как
|
||
разбор команд, вызов LLM или выполнение сервисов.
|
||
"""
|
||
|
||
from langchain_core.messages import HumanMessage, SystemMessage
|
||
|
||
from llm_client import get_llm
|
||
from services import ImageGenerationService, SubtitleService, ImageAnalysisService, SummarizationService
|
||
from models import AgentState
|
||
|
||
from typing import Dict, Any, Optional
|
||
|
||
|
||
class CommandManager:
|
||
"""
|
||
Менеджер для хранения и получения слеш-команд.
|
||
В реальном приложении команды будут храниться в БД.
|
||
"""
|
||
|
||
def __init__(self):
|
||
self._commands = {
|
||
"help": {
|
||
"description": "Показывает список доступных команд.",
|
||
"workflow": "help_workflow"
|
||
},
|
||
"imagine": {
|
||
"description":
|
||
"Генерирует изображения по промпту. Использование: /imagine [prompt]",
|
||
"workflow": "image_generation_workflow"
|
||
},
|
||
"analyze": {
|
||
"description":
|
||
"Анализирует изображение по URL и промпту. Использование: /analyze [url] [prompt]",
|
||
"workflow": "image_analysis_workflow"
|
||
},
|
||
"subtitles_meet": {
|
||
"description": "Получает текущие субтитры из Google Meet.",
|
||
"workflow": "get_meet_subtitles_workflow"
|
||
},
|
||
"subtitles_teams": {
|
||
"description": "Получает текущие субтитры из MS Teams.",
|
||
"workflow": "get_teams_subtitles_workflow"
|
||
},
|
||
"summarize": {
|
||
"description": "Суммирует текущую историю диалога.",
|
||
"workflow": "summarize_history_workflow"
|
||
},
|
||
"chat": {
|
||
"description": "Ведет обычный диалог с LLM.",
|
||
"workflow": "chat_workflow"
|
||
},
|
||
}
|
||
|
||
def get_commands(self) -> Dict[str, Any]:
|
||
"""Возвращает все доступные команды."""
|
||
return self._commands
|
||
|
||
def get_command_workflow(self, command_name: str) -> Optional[str]:
|
||
"""Возвращает имя workflow для заданной команды."""
|
||
return self._commands.get(command_name, {}).get("workflow")
|
||
|
||
|
||
# Инициализация сервисов (пример, убедитесь, что они доступны в этом файле или импортированы)
|
||
command_manager = CommandManager()
|
||
|
||
|
||
def parse_command_node(state: AgentState) -> AgentState:
|
||
"""Парсит вход пользователя на предмет слеш-команды."""
|
||
print("Выполняется узел: parse_command_node")
|
||
user_input = state.input.strip()
|
||
|
||
# Добавляем узел пользователя в граф для фронтенда
|
||
user_node_id = f"user_{len(state.graph_nodes) + 1}"
|
||
state.graph_nodes.append({
|
||
"id": user_node_id,
|
||
"type": "user",
|
||
"data": {
|
||
"label":
|
||
user_input[:30] + "..." if len(user_input) > 30 else user_input
|
||
}
|
||
})
|
||
if state.parent_node_id:
|
||
state.graph_edges.append({
|
||
"id": f"e{state.parent_node_id}-{user_node_id}",
|
||
"source": state.parent_node_id,
|
||
"target": user_node_id
|
||
})
|
||
state.parent_node_id = user_node_id
|
||
state.current_node_id = user_node_id
|
||
|
||
state.chat_history.append({
|
||
"role": "user",
|
||
"content": user_input,
|
||
"node_id": user_node_id # Сохраняем ID узла
|
||
})
|
||
|
||
if user_input.startswith('/'):
|
||
parts = user_input[1:].split(' ', 1)
|
||
command_name = parts[0].lower()
|
||
command_args = parts[1].split(' ') if len(parts) > 1 else []
|
||
if command_name in command_manager.get_commands():
|
||
state.command = command_name
|
||
state.command_args = command_args
|
||
print(
|
||
f"Команда распознана: /{state.command} с аргументами: {state.command_args}"
|
||
)
|
||
else:
|
||
state.error = f"Неизвестная команда: /{command_name}. Введите /help для списка команд."
|
||
print(state.error)
|
||
state.command = "error" # Специальный флаг для обработки ошибок команд
|
||
else:
|
||
state.command = "chat" # Обычный чат
|
||
return state
|
||
|
||
|
||
def execute_command_node(state: AgentState) -> AgentState:
|
||
"""Выполняет соответствующий workflow для команды."""
|
||
print(
|
||
f"Выполняется узел: execute_command_node для команды: {state.command}")
|
||
# Этот узел будет выступать в роли маршрутизатора
|
||
# Фактическая логика команды будет в отдельных узлах, вызываемых по условию
|
||
return state # Просто передаем состояние дальше, чтобы маршрутизатор смог выбрать следующий узел
|
||
|
||
|
||
# Предполагаем, что используем Gemini для основного LLM и суммаризации/анализа
|
||
# В реальной системе можно выбрать динамически
|
||
DEFAULT_LLM_NAME = "gemini-2.5-flash"
|
||
main_llm = get_llm(DEFAULT_LLM_NAME)
|
||
|
||
|
||
def call_llm_node(state: AgentState) -> AgentState:
|
||
"""Вызывает LLM для обработки обычного диалога или генерации ответа."""
|
||
print("Выполняется узел: call_llm_node (для обычного чата)")
|
||
try:
|
||
# Для обычного чата LLM использует всю историю
|
||
messages_for_llm = []
|
||
for msg in state.chat_history:
|
||
if msg["role"] == "user":
|
||
messages_for_llm.append(HumanMessage(content=msg["content"]))
|
||
elif msg["role"] == "assistant":
|
||
messages_for_llm.append(SystemMessage(content=msg["content"]))
|
||
|
||
response = main_llm.invoke(messages_for_llm)
|
||
|
||
# Добавляем узел LLM в граф для фронтенда
|
||
llm_node_id = f"llm_{len(state.graph_nodes) + 1}"
|
||
state.graph_nodes.append({
|
||
"id": llm_node_id,
|
||
"type": "llm",
|
||
"data": {
|
||
"label":
|
||
response[:30] + "..." if len(response) > 30 else response
|
||
}
|
||
})
|
||
state.graph_edges.append({
|
||
"id": f"e{state.parent_node_id}-{llm_node_id}",
|
||
"source": state.parent_node_id,
|
||
"target": llm_node_id
|
||
})
|
||
state.parent_node_id = llm_node_id
|
||
state.current_node_id = llm_node_id
|
||
|
||
state.chat_history.append({
|
||
"role": "assistant",
|
||
"content": response,
|
||
"node_id": llm_node_id # Сохраняем ID узла
|
||
})
|
||
state.llm_response = response
|
||
|
||
except Exception as e:
|
||
state.error = f"Ошибка при вызове LLM: {e}"
|
||
state.llm_response = f"Произошла ошибка: {e}"
|
||
print(state.error)
|
||
return state
|
||
|
||
|
||
image_gen_service = ImageGenerationService(
|
||
base_url="https://render-service-gsu7.onrender.com/g2"
|
||
) # TODO: Замените на реальный URL
|
||
|
||
|
||
def generate_images_node(state: AgentState) -> AgentState:
|
||
"""Генерирует изображения по промпту."""
|
||
print("Выполняется узел: generate_images_node")
|
||
prompt = " ".join(state.command_args)
|
||
if not prompt:
|
||
state.error = "Для команды /imagine требуется промпт. Пример: /imagine кошка на луне"
|
||
state.llm_response = state.error
|
||
return state
|
||
|
||
try:
|
||
urls = image_gen_service.generate_images(prompt)
|
||
state.image_urls = urls
|
||
|
||
# Формируем текст ответа с использованием Markdown для отображения изображений
|
||
image_markdown = "\n".join([
|
||
f'<a href="{url}"><img src="{url}" alt="image{i}" style="max-width: 300px; max-height: 300px;"></a>'
|
||
for i, url in enumerate(urls)
|
||
])
|
||
response_text = f"Сгенерировано {len(urls)} изображений:\n{image_markdown}"
|
||
state.llm_response = response_text
|
||
|
||
# Добавляем узел генерации изображений в граф
|
||
img_gen_node_id = f"img_gen_{len(state.graph_nodes) + 1}"
|
||
state.graph_nodes.append({
|
||
"id": img_gen_node_id,
|
||
"type": "tool",
|
||
"data": {
|
||
"label": "Генерация изображений",
|
||
"details": prompt
|
||
}
|
||
})
|
||
state.graph_edges.append({
|
||
"id": f"e{state.parent_node_id}-{img_gen_node_id}",
|
||
"source": state.parent_node_id,
|
||
"target": img_gen_node_id
|
||
})
|
||
state.parent_node_id = img_gen_node_id
|
||
state.current_node_id = img_gen_node_id
|
||
|
||
state.llm_response = "Сгенерированы изображения."
|
||
state.chat_history.append({
|
||
"role": "assistant",
|
||
"content": response_text,
|
||
"node_id": img_gen_node_id # Привязываем к ID узла
|
||
})
|
||
|
||
except Exception as e:
|
||
state.error = f"Ошибка генерации изображений: {e}"
|
||
state.llm_response = f"Произошла ошибка при генерации изображений: {e}"
|
||
print(state.error)
|
||
return state
|
||
|
||
|
||
image_analysis_service = ImageAnalysisService(
|
||
llm=get_llm("gemini-2.5-flash")) # Gemini для анализа изображений
|
||
|
||
|
||
def analyze_image_node(state: AgentState) -> AgentState:
|
||
"""Анализирует изображение по URL или данным."""
|
||
print("Выполняется узел: analyze_image_node")
|
||
if len(state.command_args) < 2:
|
||
state.error = "Для команды /analyze требуется URL изображения и промпт. Пример: /analyze [url] 'что на изображении?'"
|
||
state.llm_response = state.error
|
||
return state
|
||
|
||
image_url = state.command_args[0]
|
||
prompt = " ".join(state.command_args[1:])
|
||
|
||
try:
|
||
analysis_result = image_analysis_service.analyze_image(
|
||
image_url, prompt)
|
||
state.analysis_result = analysis_result
|
||
state.llm_response = f"Результат анализа изображения '{image_url}': {analysis_result}"
|
||
state.chat_history.append({
|
||
"role": "assistant",
|
||
"content": state.llm_response
|
||
})
|
||
|
||
# Добавляем узел анализа изображения в граф
|
||
img_analysis_node_id = f"img_analysis_{len(state.graph_nodes) + 1}"
|
||
state.graph_nodes.append({
|
||
"id": img_analysis_node_id,
|
||
"type": "tool",
|
||
"data": {
|
||
"label": "Анализ изображения",
|
||
"details": f"URL: {image_url}, Промпт: {prompt}"
|
||
}
|
||
})
|
||
state.graph_edges.append({
|
||
"id": f"e{state.parent_node_id}-{img_analysis_node_id}",
|
||
"source": state.parent_node_id,
|
||
"target": img_analysis_node_id
|
||
})
|
||
state.parent_node_id = img_analysis_node_id
|
||
state.current_node_id = img_analysis_node_id
|
||
|
||
except Exception as e:
|
||
state.error = f"Ошибка анализа изображения: {e}"
|
||
state.llm_response = f"Произошла ошибка при анализе изображения: {e}"
|
||
print(state.error)
|
||
return state
|
||
|
||
|
||
subtitle_service = SubtitleService()
|
||
|
||
|
||
def get_meet_subtitles_node(state: AgentState) -> AgentState:
|
||
"""Получает субтитры из Google Meet."""
|
||
print("Выполняется узел: get_meet_subtitles_node")
|
||
try:
|
||
subtitles = subtitle_service.get_google_meet_subtitles()
|
||
state.subtitles = subtitles
|
||
state.llm_response = f"Текущие субтитры из Google Meet: {subtitles}"
|
||
state.chat_history.append({
|
||
"role": "assistant",
|
||
"content": state.llm_response
|
||
})
|
||
|
||
# Добавляем узел субтитров Meet в граф
|
||
meet_subtitles_node_id = f"meet_subtitles_{len(state.graph_nodes) + 1}"
|
||
state.graph_nodes.append({
|
||
"id": meet_subtitles_node_id,
|
||
"type": "tool",
|
||
"data": {
|
||
"label": "Субтитры Google Meet",
|
||
"details": subtitles[:30] + "..." if subtitles else ""
|
||
}
|
||
})
|
||
state.graph_edges.append({
|
||
"id": f"e{state.parent_node_id}-{meet_subtitles_node_id}",
|
||
"source": state.parent_node_id,
|
||
"target": meet_subtitles_node_id
|
||
})
|
||
state.parent_node_id = meet_subtitles_node_id
|
||
state.current_node_id = meet_subtitles_node_id
|
||
|
||
except Exception as e:
|
||
state.error = f"Ошибка получения субтитров Google Meet: {e}"
|
||
state.llm_response = f"Произошла ошибка при получении субтитров Google Meet: {e}"
|
||
print(state.error)
|
||
return state
|
||
|
||
|
||
def get_teams_subtitles_node(state: AgentState) -> AgentState:
|
||
"""Получает субтитры из MS Teams."""
|
||
print("Выполняется узел: get_teams_subtitles_node")
|
||
try:
|
||
subtitles = subtitle_service.get_ms_teams_subtitles()
|
||
state.subtitles = subtitles
|
||
state.llm_response = f"Текущие субтитры из MS Teams: {subtitles}"
|
||
state.chat_history.append({
|
||
"role": "assistant",
|
||
"content": state.llm_response
|
||
})
|
||
|
||
# Добавляем узел субтитров Teams в граф
|
||
teams_subtitles_node_id = f"teams_subtitles_{len(state.graph_nodes) + 1}"
|
||
state.graph_nodes.append({
|
||
"id": teams_subtitles_node_id,
|
||
"type": "tool",
|
||
"data": {
|
||
"label": "Субтитры MS Teams",
|
||
"details": subtitles[:30] + "..." if subtitles else ""
|
||
}
|
||
})
|
||
state.graph_edges.append({
|
||
"id": f"e{state.parent_node_id}-{teams_subtitles_node_id}",
|
||
"source": state.parent_node_id,
|
||
"target": teams_subtitles_node_id
|
||
})
|
||
state.parent_node_id = teams_subtitles_node_id
|
||
state.current_node_id = teams_subtitles_node_id
|
||
|
||
except Exception as e:
|
||
state.error = f"Ошибка получения субтитров MS Teams: {e}"
|
||
state.llm_response = f"Произошла ошибка при получении субтитров MS Teams: {e}"
|
||
print(state.error)
|
||
return state
|
||
|
||
|
||
summarization_service = SummarizationService(
|
||
llm=get_llm("gemini-2.5-flash")) # Gemini для суммаризации
|
||
|
||
|
||
def summarize_history_node(state: AgentState) -> AgentState:
|
||
"""Суммирует историю диалога."""
|
||
print("Выполняется узел: summarize_history_node")
|
||
try:
|
||
# Для суммаризации используем полную историю, которую ведет AgentState
|
||
summarized_text = summarization_service.summarize_history(
|
||
state.chat_history)
|
||
state.summarized_history_text = summarized_text
|
||
state.llm_response = f"История диалога суммирована:\n{summarized_text}"
|
||
state.chat_history.append({
|
||
"role": "assistant",
|
||
"content": state.llm_response
|
||
})
|
||
|
||
# Добавляем узел суммаризации в граф
|
||
summary_node_id = f"summary_{len(state.graph_nodes) + 1}"
|
||
state.graph_nodes.append({
|
||
"id": summary_node_id,
|
||
"type": "tool",
|
||
"data": {
|
||
"label": "Суммаризация истории",
|
||
"details":
|
||
summarized_text[:30] + "..." if summarized_text else ""
|
||
}
|
||
})
|
||
state.graph_edges.append({
|
||
"id": f"e{state.parent_node_id}-{summary_node_id}",
|
||
"source": state.parent_node_id,
|
||
"target": summary_node_id
|
||
})
|
||
state.parent_node_id = summary_node_id
|
||
state.current_node_id = summary_node_id
|
||
|
||
except Exception as e:
|
||
state.error = f"Ошибка суммаризации истории: {e}"
|
||
state.llm_response = f"Произошла ошибка при суммаризации истории: {e}"
|
||
print(state.error)
|
||
return state
|
||
|
||
|
||
def handle_error_node(state: AgentState) -> AgentState:
|
||
"""Обрабатывает ошибки команд."""
|
||
print(f"Выполняется узел: handle_error_node. Ошибка: {state.error}")
|
||
state.llm_response = state.error
|
||
state.chat_history.append({
|
||
"role": "assistant",
|
||
"content": state.llm_response
|
||
})
|
||
|
||
# Добавляем узел ошибки в граф
|
||
error_node_id = f"error_{len(state.graph_nodes) + 1}"
|
||
state.graph_nodes.append({
|
||
"id": error_node_id,
|
||
"type": "error",
|
||
"data": {
|
||
"label": "Ошибка",
|
||
"details": state.error
|
||
}
|
||
})
|
||
state.graph_edges.append({
|
||
"id": f"e{state.parent_node_id}-{error_node_id}",
|
||
"source": state.parent_node_id,
|
||
"target": error_node_id
|
||
})
|
||
state.parent_node_id = error_node_id
|
||
state.current_node_id = error_node_id
|
||
|
||
return state
|
||
|
||
|
||
def help_node(state: AgentState) -> AgentState:
|
||
"""Предоставляет информацию о доступных командах."""
|
||
print("Выполняется узел: help_node")
|
||
commands_info = command_manager.get_commands()
|
||
help_text = "Доступные команды:\n"
|
||
for cmd, info in commands_info.items():
|
||
help_text += f"/{cmd}: {info['description']}\n"
|
||
state.llm_response = help_text
|
||
state.chat_history.append({
|
||
"role": "assistant",
|
||
"content": state.llm_response
|
||
})
|
||
|
||
# Добавляем узел справки в граф
|
||
help_node_id = f"help_{len(state.graph_nodes) + 1}"
|
||
state.graph_nodes.append({
|
||
"id": help_node_id,
|
||
"type": "info",
|
||
"data": {
|
||
"label": "Справка",
|
||
"details": "Список команд"
|
||
}
|
||
})
|
||
state.graph_edges.append({
|
||
"id": f"e{state.parent_node_id}-{help_node_id}",
|
||
"source": state.parent_node_id,
|
||
"target": help_node_id
|
||
})
|
||
state.parent_node_id = help_node_id
|
||
state.current_node_id = help_node_id
|
||
|
||
return state
|