""" Асинхронный сервис для генерации заголовков графов и узлов. Работает в фоновом режиме, обрабатывая очередь элементов без заголовков. """ import sqlite3 import threading import time from typing import Any, Dict, Optional 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, db_path="graph_history.db"): self.db_path = db_path self._create_table() self.history_manager = history_manager self.llm = get_llm(DEFAULT_SUMMARIZATION_LLM_NAME) self.running = False self.thread: Optional[threading.Thread] = None self.socketio = None # Системные промпты для генерации заголовков self.graph_title_prompt = """Создай краткий заголовок (максимум 80 символов) для диалога на основе первого сообщения пользователя. Заголовок должен отражать основную тему или вопрос. Заголовок должен быть простым текстом, без какого-либо форматирования или использования специальных символов разметки (например, #, *, _, `). Отвечай только заголовком, без дополнительных объяснений. Ты должен сделать саммери, а не ответить на вопросы, если они есть в сообщении.""" self.node_title_prompt = """Создай краткий заголовок (максимум 40 символов) для этого сообщения/действия. Заголовок должен кратко описывать суть сообщения или действия. Заголовок должен быть простым текстом, без какого-либо форматирования или использования специальных символов разметки (например, #, *, _, `). Отвечай только заголовком, без дополнительных объяснений. Ты должен сделать саммери, а не ответить на вопросы, если они есть в сообщении.""" def _get_connection(self): """ Получает соединение с базой данных. Устанавливаем таймаут для ожидания блокировки базы данных при конкурентной записи. """ return sqlite3.connect(self.db_path, timeout=5.0) # Увеличил таймаут до 5 секунд def _create_table(self): """Создает таблицы для хранения графов, если они не существуют.""" with self._get_connection() as conn: cursor = conn.cursor() # Очередь для обработки заголовков cursor.execute(""" CREATE TABLE IF NOT EXISTS title_generation_queue ( id INTEGER PRIMARY KEY AUTOINCREMENT, item_type TEXT NOT NULL, -- 'graph' or 'node' graph_id TEXT NOT NULL, node_id TEXT DEFAULT NULL, -- только для узлов priority INTEGER DEFAULT 0, -- для приоритета обработки created_at DATETIME DEFAULT CURRENT_TIMESTAMP ) """) conn.commit() def start(self): """Запускает фоновый процесс генерации заголовков.""" if self.running: return self.running = True # Заполняем очередь при старте self.populate_initial_title_queue() # Запускаем фоновый поток self.thread = threading.Thread(target=self._process_queue, daemon=True) self.thread.start() print("Сервис генерации заголовков запущен") def populate_initial_title_queue(self): """Заполняет очередь элементами без сгенерированных заголовков при старте сервера.""" with self._get_connection() as conn: cursor = conn.cursor() # Очищаем существующую очередь cursor.execute("DELETE FROM title_generation_queue") # Добавляем графы без заголовков cursor.execute( "SELECT id FROM graphs WHERE title_generated = FALSE OR title_generated IS NULL" ) for (graph_id, ) in cursor.fetchall(): cursor.execute( "INSERT INTO title_generation_queue (item_type, graph_id, priority) VALUES (?, ?, ?)", ("graph", graph_id, 1)) # Добавляем узлы без заголовков cursor.execute( "SELECT graph_id, node_id FROM graph_nodes_data WHERE title_generated = FALSE OR title_generated IS NULL" ) for graph_id, node_id in cursor.fetchall(): cursor.execute( "INSERT INTO title_generation_queue (item_type, graph_id, node_id, priority) VALUES (?, ?, ?, ?)", ("node", graph_id, node_id, 1)) conn.commit() def stop(self): """Останавливает фоновый процесс.""" self.running = False if self.thread: self.thread.join() print("Сервис генерации заголовков остановлен") def _process_queue(self): """Основной цикл обработки очереди заголовков.""" while self.running: try: item = self.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 get_next_from_title_queue(self) -> Optional[Dict[str, Any]]: """Получает следующий элемент из очереди для обработки.""" with self._get_connection() as conn: cursor = conn.cursor() cursor.execute(""" SELECT id, item_type, graph_id, node_id FROM title_generation_queue ORDER BY priority DESC, id ASC LIMIT 1 """) result = cursor.fetchone() if result: queue_id, item_type, graph_id, node_id = result # Удаляем из очереди cursor.execute( "DELETE FROM title_generation_queue WHERE id = ?", (queue_id, )) conn.commit() return { "item_type": item_type, "graph_id": graph_id, "node_id": node_id } return None 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}") # ----------------------------------------------- WebSocket Notification --------------------------------------------------------------- # Отправляем событие через WebSocket if self.socketio: self.socketio.emit('graph_title_updated', { 'graph_id': graph_id, 'title': 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}") # ----------------------------------------------- WebSocket Notification --------------------------------------------------------------- # Отправляем событие через WebSocket if self.socketio: self.socketio.emit('node_title_updated', { 'graph_id': graph_id, 'node_id': node_id, 'title': title }) except Exception as e: print(f"Ошибка генерации заголовка узла {node_id}: {e}") def add_node_to_queue(self, cursor, graph_id, node_id): # Добавляем в очередь заголовков с высоким приоритетом cursor.execute( "INSERT INTO title_generation_queue (item_type, graph_id, node_id, priority) VALUES (?, ?, ?, ?)", ("node", graph_id, node_id, 10)) def add_graph_to_queue(self, cursor, graph_id): cursor.execute( "INSERT INTO title_generation_queue (item_type, graph_id, priority) VALUES (?, ?, ?)", ("graph", graph_id, 10)) def add_node_to_queue_direct(self, graph_id: str, node_id: str, priority: int = 10): """Добавляет узел в очередь генерации заголовков напрямую (с собственным соединением).""" with self._get_connection() as conn: cursor = conn.cursor() cursor.execute( "INSERT INTO title_generation_queue (item_type, graph_id, node_id, priority) VALUES (?, ?, ?, ?)", ("node", graph_id, node_id, priority)) conn.commit() def add_graph_to_queue_direct(self, graph_id: str, priority: int = 10): """Добавляет граф в очередь генерации заголовков напрямую (с собственным соединением).""" with self._get_connection() as conn: cursor = conn.cursor() cursor.execute( "INSERT INTO title_generation_queue (item_type, graph_id, priority) VALUES (?, ?, ?)", ("graph", graph_id, priority)) conn.commit()