151 lines
7.1 KiB
Python
151 lines
7.1 KiB
Python
"""
|
||
Асинхронный сервис для генерации заголовков графов и узлов.
|
||
Работает в фоновом режиме, обрабатывая очередь элементов без заголовков.
|
||
"""
|
||
|
||
import threading
|
||
import time
|
||
from typing import Optional
|
||
from graph_history_manager import GraphHistoryManager
|
||
from nodes import DEFAULT_SUMMARIZATION_LLM_NAME
|
||
from llm_client import get_llm
|
||
from langchain_core.messages import SystemMessage, HumanMessage
|
||
|
||
class TitleGenerator:
|
||
"""Сервис для асинхронной генерации заголовков."""
|
||
|
||
def __init__(self, history_manager: GraphHistoryManager):
|
||
self.history_manager = history_manager
|
||
self.llm = get_llm(DEFAULT_SUMMARIZATION_LLM_NAME)
|
||
self.running = False
|
||
self.thread: Optional[threading.Thread] = None
|
||
|
||
# Системные промпты для генерации заголовков
|
||
self.graph_title_prompt = """Создай краткий заголовок (максимум 80 символов) для диалога на основе первого сообщения пользователя.
|
||
Заголовок должен отражать основную тему или вопрос. Отвечай только заголовком, без дополнительных объяснений. Ты должен сделать саммери, а не ответить на вопросы, если они есть в сообщении."""
|
||
|
||
self.node_title_prompt = """Создай краткий заголовок (максимум 40 символов) для этого сообщения/действия.
|
||
Заголовок должен кратко описывать суть сообщения или действия. Отвечай только заголовком, без дополнительных объяснений. Ты должен сделать саммери, а не ответить на вопросы, если они есть в сообщении."""
|
||
|
||
def start(self):
|
||
"""Запускает фоновый процесс генерации заголовков."""
|
||
if self.running:
|
||
return
|
||
|
||
self.running = True
|
||
|
||
# Заполняем очередь при старте
|
||
self.history_manager.populate_initial_title_queue()
|
||
|
||
# Запускаем фоновый поток
|
||
self.thread = threading.Thread(target=self._process_queue, daemon=True)
|
||
self.thread.start()
|
||
print("Сервис генерации заголовков запущен")
|
||
|
||
def stop(self):
|
||
"""Останавливает фоновый процесс."""
|
||
self.running = False
|
||
if self.thread:
|
||
self.thread.join()
|
||
print("Сервис генерации заголовков остановлен")
|
||
|
||
def _process_queue(self):
|
||
"""Основной цикл обработки очереди заголовков."""
|
||
while self.running:
|
||
try:
|
||
item = self.history_manager.get_next_from_title_queue()
|
||
|
||
if item:
|
||
if item["item_type"] == "graph":
|
||
self._generate_graph_title(item["graph_id"])
|
||
elif item["item_type"] == "node":
|
||
self._generate_node_title(item["graph_id"], item["node_id"])
|
||
else:
|
||
# Если очередь пуста, ждем немного
|
||
time.sleep(5)
|
||
|
||
except Exception as e:
|
||
print(f"Ошибка при обработке очереди заголовков: {e}")
|
||
time.sleep(10)
|
||
|
||
def _generate_graph_title(self, graph_id: str):
|
||
"""Генерирует заголовок для графа."""
|
||
try:
|
||
graph_data = self.history_manager.get_graph(graph_id)
|
||
if not graph_data:
|
||
return
|
||
|
||
messages = graph_data.get("messages", [])
|
||
if not messages:
|
||
return
|
||
|
||
# Находим первое пользовательское сообщение
|
||
first_user_message = None
|
||
for msg in messages:
|
||
if msg.get("role") == "user":
|
||
first_user_message = msg.get("content", "")
|
||
break
|
||
|
||
if not first_user_message:
|
||
return
|
||
|
||
# Генерируем заголовок
|
||
llm_messages = [
|
||
SystemMessage(content=self.graph_title_prompt),
|
||
HumanMessage(content=first_user_message)
|
||
]
|
||
|
||
title = self.llm.invoke(llm_messages)
|
||
title = title.strip()[:80] # Ограничиваем длину
|
||
|
||
# Сохраняем заголовок
|
||
self.history_manager.update_graph_title(graph_id, title)
|
||
print(f"Сгенерирован заголовок графа {graph_id}: {title}")
|
||
|
||
except Exception as e:
|
||
print(f"Ошибка генерации заголовка графа {graph_id}: {e}")
|
||
|
||
def _generate_node_title(self, graph_id: str, node_id: str):
|
||
"""Генерирует заголовок для узла."""
|
||
try:
|
||
graph_data = self.history_manager.get_graph(graph_id)
|
||
if not graph_data:
|
||
return
|
||
|
||
# Находим узел по ID
|
||
target_node = None
|
||
for node in graph_data.get("graph_nodes", []):
|
||
if node["id"] == node_id:
|
||
target_node = node
|
||
break
|
||
|
||
if not target_node:
|
||
return
|
||
|
||
# Получаем содержимое узла
|
||
node_data = target_node.get("data", {})
|
||
message = node_data.get("message", {})
|
||
content = message.get("content", "")
|
||
|
||
if not content:
|
||
# Если нет content, используем label или type
|
||
content = node_data.get("label", target_node.get("type", ""))
|
||
|
||
if not content:
|
||
return
|
||
|
||
# Генерируем заголовок
|
||
llm_messages = [
|
||
SystemMessage(content=self.node_title_prompt),
|
||
HumanMessage(content=content)
|
||
]
|
||
|
||
title = self.llm.invoke(llm_messages)
|
||
title = title.strip()[:40] # Ограничиваем длину
|
||
|
||
# Сохраняем заголовок
|
||
self.history_manager.update_node_title(graph_id, node_id, title)
|
||
print(f"Сгенерирован заголовок узла {node_id}: {title}")
|
||
|
||
except Exception as e:
|
||
print(f"Ошибка генерации заголовка узла {node_id}: {e}") |