# отключаем 40-ка строчное ограничение для этого файла # ----------------------------------------------- Imports ------------------------------------------------------------ from typing import Dict, Any, List, Optional import json import sqlite3 import uuid # Для генерации UUID from title_generator import TitleGenerator # ----------------------------------------------- Exceptions ------------------------------------------------------------ class GraphUpdateConflictError(Exception): """Исключение возникает при конфликте обновления графа.""" pass # ----------------------------------------------- GraphHistoryManager ------------------------------------------------------------ class GraphHistoryManager: """ Менеджер для хранения истории графов запросов в SQLite. Использует отдельные таблицы для узлов и ребер для инкрементального обновления. Обеспечивает конкурентный доступ к разным графам, полагаясь на транзакционность SQLite. """ def __init__(self, db_path="graph_history.db"): self.db_path = db_path self._create_tables() # Инициализируем сервис генерации заголовков self.title_generator = TitleGenerator(self) self.title_generator.start() self.socketio = None def set_socketio(self, socketio): """ Устанавливает объект SocketIO для отправки уведомлений генератору заголовков. @param socketio: Объект SocketIO. """ self.socketio = socketio self.title_generator.set_socketio(socketio) def __del__(self): """ Деструктор класса. Останавливает генератор заголовков при уничтожении объекта. """ print("🚩 Завершение работы GraphHistoryManager. Останавливаем генератор заголовков...") self.stop_title_generator() def stop_title_generator(self): """Останавливает процесс генерации заголовков.""" if self.title_generator: self.title_generator.stop() # ----------------------------------------------- Internal Methods ------------------------------------------------------------ def _get_connection(self): """ Получает соединение с базой данных. Устанавливаем таймаут для ожидания блокировки базы данных при конкурентной записи. """ return sqlite3.connect(self.db_path, timeout=5.0) # Увеличил таймаут до 5 секунд def _create_tables(self): """Создает таблицы для хранения графов, если они не существуют.""" with self._get_connection() as conn: cursor = conn.cursor() cursor.execute(""" CREATE TABLE IF NOT EXISTS graphs ( id TEXT PRIMARY KEY, current_node_id TEXT DEFAULT NULL, title TEXT DEFAULT NULL, title_generated BOOLEAN DEFAULT FALSE, custom_system_prompt TEXT DEFAULT NULL, timestamp DATETIME DEFAULT CURRENT_TIMESTAMP ) """) # временно так cursor.execute("PRAGMA table_info(graphs)") columns = [col[1] for col in cursor.fetchall()] if 'custom_system_prompt' not in columns: cursor.execute( "ALTER TABLE graphs ADD COLUMN custom_system_prompt TEXT DEFAULT NULL" ) cursor.execute(""" CREATE TABLE IF NOT EXISTS graph_nodes_data ( graph_id TEXT NOT NULL, node_id TEXT NOT NULL, node_type TEXT NOT NULL, node_data TEXT NOT NULL, title TEXT DEFAULT NULL, title_generated BOOLEAN DEFAULT FALSE, PRIMARY KEY (graph_id, node_id), FOREIGN KEY (graph_id) REFERENCES graphs (id) ON DELETE CASCADE ) """) cursor.execute(""" CREATE TABLE IF NOT EXISTS graph_edges_data ( graph_id TEXT NOT NULL, edge_id TEXT NOT NULL, source_node_id TEXT NOT NULL, target_node_id TEXT NOT NULL, PRIMARY KEY (graph_id, edge_id), FOREIGN KEY (graph_id) REFERENCES graphs (id) ON DELETE CASCADE, FOREIGN KEY (source_node_id) REFERENCES graph_nodes_data (node_id) ON DELETE CASCADE, FOREIGN KEY (target_node_id) REFERENCES graph_nodes_data (node_id) ON DELETE CASCADE ) """) conn.commit() def _ensure_graph_exists(self, graph_id: str) -> bool: """ Проверяет существование графа с заданным ID. @param graph_id: ID графа. @returns: True, если граф существует, False в противном случае. """ with self._get_connection() as conn: cursor = conn.cursor() cursor.execute("SELECT 1 FROM graphs WHERE id = ?", (graph_id,)) return cursor.fetchone() is not None # ----------------------------------------------- Graph Node/Edge Access Helpers ------------------------------------------------------------ def _get_graph_nodes(self, conn, graph_id: str) -> List[Dict[str, Any]]: """Извлекает все узлы для заданного графа.""" cursor = conn.cursor() cursor.execute( "SELECT node_id, node_type, node_data, title, title_generated FROM graph_nodes_data WHERE graph_id = ?", (graph_id, )) nodes = [] for node_id, node_type, node_data_json, title, title_generated in cursor.fetchall(): node_data = json.loads(node_data_json) # Добавляем заголовок в данные узла if title and title_generated: node_data["title"] = title node_data["title_generated"] = bool(title_generated) nodes.append({ "id": node_id, "type": node_type, "data": node_data }) return nodes def _get_graph_edges(self, conn, graph_id: str) -> List[Dict[str, Any]]: """Извлекает все ребра для заданного графа.""" cursor = conn.cursor() cursor.execute( "SELECT edge_id, source_node_id, target_node_id FROM graph_edges_data WHERE graph_id = ?", (graph_id, )) edges = [] for edge_id, source, target in cursor.fetchall(): edges.append({"id": edge_id, "source": source, "target": target}) return edges def _update_graph_current_node_id(self, conn, graph_id: str, new_current_node_id: str): """Обновляет current_node_id для графа.""" cursor = conn.cursor() cursor.execute("UPDATE graphs SET current_node_id = ? WHERE id = ?", (new_current_node_id, graph_id)) def _add_graph_node(self, conn, graph_id: str, node: Dict[str, Any]): """Добавляет новый узел в граф. Использует INSERT OR IGNORE для избежания конфликтов.""" cursor = conn.cursor() cursor.execute( "INSERT OR IGNORE INTO graph_nodes_data (graph_id, node_id, node_type, node_data) VALUES (?, ?, ?, ?)", (graph_id, node["id"], node["type"], json.dumps(node["data"]))) def _add_graph_edge(self, conn, graph_id: str, edge: Dict[str, Any]): """Добавляет новое ребро в граф. Использует INSERT OR IGNORE для избежания конфликтов.""" cursor = conn.cursor() cursor.execute( "INSERT OR IGNORE INTO graph_edges_data (graph_id, edge_id, source_node_id, target_node_id) VALUES (?, ?, ?, ?)", (graph_id, edge["id"], edge["source"], edge["target"])) # ----------------------------------------------- Public Methods ------------------------------------------------------------ def create_new_graph(self) -> str: """ Создает новый пустой граф в базе данных и возвращает его ID. @returns: ID нового графа. """ new_graph_id = str(uuid.uuid4()) with self._get_connection() as conn: cursor = conn.cursor() try: cursor.execute( "INSERT INTO graphs (id, title, title_generated, custom_system_prompt) VALUES (?, ?, FALSE, NULL)", (new_graph_id, "Новый граф") # Устанавливаем заголовок по умолчанию ) conn.commit() return new_graph_id except Exception as e: conn.rollback() print(f"Ошибка при создании нового графа: {e}") raise def get_graph( self, graph_id: str, target_node_id: Optional[str] = None) -> Optional[Dict[str, Any]]: """ Получает данные графа по ID из базы данных. """ if not self._ensure_graph_exists(graph_id): print(f"Граф с ID {graph_id} не найден.") return None with self._get_connection() as conn: cursor = conn.cursor() cursor.execute("SELECT current_node_id, title, title_generated, custom_system_prompt FROM graphs WHERE id = ?", # Извлекаем title и title_generated (graph_id, )) result = cursor.fetchone() if result: db_current_node_id, graph_title, graph_title_generated, custom_system_prompt = result # Получаем все узлы и ребра graph_nodes = self._get_graph_nodes(conn, graph_id) graph_edges = self._get_graph_edges(conn, graph_id) resolved_current_node_id = target_node_id if target_node_id else db_current_node_id messages = self.get_messages_from_root_to_node( graph_id, { "graph_nodes": graph_nodes, "graph_edges": graph_edges }, resolved_current_node_id) # Определяем "первое_сообщение" для отображения, если заголовок еще не сгенерирован first_message_content = None if messages: for msg in messages: if msg.get("role") == "user": first_message_content = msg.get("content") break if not first_message_content and messages: first_message_content = messages[0].get("content") return { "id": graph_id, "messages": messages, "graph_nodes": graph_nodes, "graph_edges": graph_edges, "current_node_id": resolved_current_node_id, "title": graph_title if graph_title_generated else None, "first_message": first_message_content, "custom_system_prompt": custom_system_prompt } else: return None def get_all_graphs_summary(self) -> List[Dict[str, str]]: """ Возвращает краткий список всех сохраненных графов, отсортированных по дате добавления. """ with self._get_connection() as conn: cursor = conn.cursor() # Сортируем по timestamp по убыванию, чтобы самые новые были сверху cursor.execute( "SELECT id, current_node_id, title, title_generated FROM graphs ORDER BY timestamp DESC" ) graphs_data = cursor.fetchall() summaries = [] for gid, current_node_id, title, title_generated in graphs_data: if title and title_generated: display_title = title else: # Fallback к первому сообщению graph_nodes = self._get_graph_nodes(conn, gid) graph_edges = self._get_graph_edges(conn, gid) # Получаем все сообщения до current_node_id full_messages = self.get_messages_from_root_to_node( gid, { "graph_nodes": graph_nodes, "graph_edges": graph_edges }, current_node_id) first_message_content = "No content" if full_messages: for msg in full_messages: if msg.get("role") == "user": first_message_content = msg.get( "content", "No content") break if first_message_content == "No content" and full_messages: first_message_content = full_messages[0].get( "content", "No content") # Обрезаем до 50 символов для отображения display_title = first_message_content[:50] + "..." if len( first_message_content) > 50 else first_message_content summaries.append({ "id": gid, "title": display_title, "title_generated": bool(title_generated) }) return summaries def get_messages_from_root_to_node( self, graph_id: str, # НОВЫЙ АРГУМЕНТ graph_data: Dict[str, Any], target_node_id: str) -> List[Dict[str, str]]: """ Получает список сообщений от корневого узла до указанного узла. Сообщения извлекаются из данных узлов, а не из отдельного поля. """ all_nodes = graph_data.get("graph_nodes", []) edges = graph_data.get("graph_edges", []) # Создаем словарь для быстрого поиска узлов по ID node_map = {node['id']: node for node in all_nodes} # Создаем словарь для быстрого поиска родительских узлов parent_map = {edge['target']: edge['source'] for edge in edges} def get_path_to_root(node_id: str) -> List[str]: path = [node_id] current = node_id while current in parent_map: current = parent_map[current] # Проверяем, чтобы не зациклиться (хотя в древовидном графе это не должно быть проблемой) if current in path: print(f"Обнаружен цикл при построении пути к узлу {node_id}. Путь: {path}") break path.append(current) return path[::-1] # Разворачиваем путь от корня до целевого узла path_to_target = get_path_to_root(target_node_id) # Собираем сообщения, соответствующие узлам в пути. # Сообщения теперь хранятся в 'data' каждого узла при создании. messages_for_path = [] for node_id in path_to_target: node = node_map.get(node_id) if node and 'message' in node.get( 'data', {}): # Проверяем наличие поля 'message' original_message = node['data']['message'] messages_for_path.append({ "role": original_message.get("role", "unknown"), "type": node.get("type", "unknown"), "content": original_message.get("content", ""), "node_id": original_message.get("node_id", node_id), "graph_id": graph_id }) return messages_for_path def delete_graph(self, graph_id: str) -> bool: """ Удаляет граф из базы данных. Благодаря CASCADE, удалит также все узлы и ребра, связанные с этим графом. """ with self._get_connection() as conn: cursor = conn.cursor() try: cursor.execute("DELETE FROM graphs WHERE id = ?", (graph_id, )) conn.commit() return cursor.rowcount > 0 except Exception as e: conn.rollback() print(f"Ошибка при удалении графа: {e}") return False def delete_all_graphs(self) -> bool: # ДОБАВЛЕНО """Удаляет все графы из базы данных.""" with self._get_connection() as conn: cursor = conn.cursor() try: cursor.execute("DELETE FROM graphs") conn.commit() return True except Exception as e: conn.rollback() print(f"Ошибка при удалении всех графов: {e}") return False def rename_graph(self, graph_id: str, new_title: str) -> bool: # ДОБАВЛЕНО """Переименовывает граф, устанавливая пользовательское название.""" with self._get_connection() as conn: cursor = conn.cursor() try: cursor.execute( "UPDATE graphs SET title = ?, title_generated = TRUE WHERE id = ?", (new_title, graph_id)) conn.commit() return cursor.rowcount > 0 except Exception as e: conn.rollback() print(f"Ошибка при переименовании графа: {e}") return False def delete_node(self, graph_id: str, node_id: str) -> bool: """ Удаляет узел из графа и переподключает его родителей к его детям. Обновляет current_node_id, если удаляемый узел был текущим. """ with self._get_connection() as conn: cursor = conn.cursor() try: # 1. Получаем информацию о текущем узле графа cursor.execute("SELECT current_node_id FROM graphs WHERE id = ?", (graph_id,)) current_graph_node_id = cursor.fetchone()[0] if cursor.rowcount > 0 else None # 2. Находим родителей и детей удаляемого узла # Родители cursor.execute("SELECT source_node_id FROM graph_edges_data WHERE graph_id = ? AND target_node_id = ?", (graph_id, node_id)) parents = [row[0] for row in cursor.fetchall()] # Дети cursor.execute("SELECT target_node_id FROM graph_edges_data WHERE graph_id = ? AND source_node_id = ?", (graph_id, node_id)) children = [row[0] for row in cursor.fetchall()] # 3. Переподключаем родителей к детям for parent_id in parents: for child_id in children: # Создаем новое ребро между родителем и ребенком new_edge_id = str(uuid.uuid4()) cursor.execute( "INSERT OR IGNORE INTO graph_edges_data (graph_id, edge_id, source_node_id, target_node_id) VALUES (?, ?, ?, ?)", (graph_id, new_edge_id, parent_id, child_id)) # 4. Удаляем все ребра, связанные с удаляемым узлом (входящие и исходящие) cursor.execute("DELETE FROM graph_edges_data WHERE graph_id = ? AND (source_node_id = ? OR target_node_id = ?)", (graph_id, node_id, node_id)) # 5. Удаляем сам узел cursor.execute("DELETE FROM graph_nodes_data WHERE graph_id = ? AND node_id = ?", (graph_id, node_id)) node_deleted = cursor.rowcount > 0 # TODO: Здесь не полная логика, т.к. если родителей нет, то current_node_id=None останется # 6. Обновляем current_node_id, если удаленный узел был текущим if node_id == current_graph_node_id: new_current_node_id = None if parents: # Если есть родители, делаем текущим узлом одного из родителей new_current_node_id = parents[0] elif children: # Если нет родителей, но есть дети, делаем текущим узлом одного из детей new_current_node_id = children[0] # Если нет ни родителей, ни детей, current_node_id останется None. cursor.execute("UPDATE graphs SET current_node_id = ? WHERE id = ?", (new_current_node_id, graph_id)) print(f"current_node_id для графа {graph_id} обновлен на {new_current_node_id}") conn.commit() return node_deleted except Exception as e: conn.rollback() print(f"Ошибка при удалении узла {node_id} из графа {graph_id}: {e}") raise # Пробрасываем исключение, чтобы API мог вернуть 500 # ----------------------------------------------- Title Update Methods ------------------------------------------------------------ def update_graph_title(self, graph_id: str, title: str): """Обновляет заголовок графа и устанавливает флаг генерации.""" with self._get_connection() as conn: cursor = conn.cursor() cursor.execute( "UPDATE graphs SET title = ?, title_generated = TRUE WHERE id = ?", (title, graph_id)) conn.commit() def update_node_title(self, graph_id: str, node_id: str, title: str): """Обновляет заголовок узла и устанавливает флаг генерации.""" with self._get_connection() as conn: cursor = conn.cursor() # Сначала проверим, существует ли узел cursor.execute( "SELECT COUNT(*) FROM graph_nodes_data WHERE graph_id = ? AND node_id = ?", (graph_id, node_id)) count = cursor.fetchone()[0] print(f"Найдено узлов для обновления: {count}") cursor.execute( "UPDATE graph_nodes_data SET title = ?, title_generated = TRUE WHERE graph_id = ? AND node_id = ?", (title, graph_id, node_id)) updated_rows = cursor.rowcount print(f"Обновлено строк: {updated_rows}") conn.commit() def create_user_node(self, message: str, graph_id: Optional[str] = None, parent_node_id: Optional[str] = None) -> Dict[str, Any]: """ Создаёт узел пользователя и возвращает graph_id и node_id. """ if not graph_id: graph_id = str(uuid.uuid4()) if not self._ensure_graph_exists(graph_id): raise Exception(f"Граф с ID {graph_id} не найден.") user_node_id = str(uuid.uuid4()) with self._get_connection() as conn: try: node_data = { "label": message[:30] + "..." if len(message) > 30 else message, "message": {"role": "user", "content": message, "node_id": user_node_id} } self._add_graph_node(conn, graph_id, { "id": user_node_id, "type": "user", "data": node_data }) if parent_node_id: edge_id = str(uuid.uuid4()) self._add_graph_edge(conn, graph_id, { "id": edge_id, "source": parent_node_id, "target": user_node_id }) self._update_graph_current_node_id(conn, graph_id, user_node_id) conn.commit() if not self.get_graph_title_generation_status(graph_id): self.title_generator.add_graph_to_queue_direct(graph_id) return { "graph_id": graph_id, "node_id": user_node_id, "parent_node_id": parent_node_id } except Exception as e: conn.rollback() print(f"Ошибка при создании узла пользователя: {e}") raise def create_assistant_placeholder_node(self, graph_id: str, parent_node_id: str) -> str: """ Создаёт узел-заглушку для ответа ассистента. """ assistant_node_id = str(uuid.uuid4()) with self._get_connection() as conn: try: node_data = { "label": "Получение ответа...", "message": {"role": "assistant", "content": "", "node_id": assistant_node_id}, "is_placeholder": True } self._add_graph_node(conn, graph_id, { "id": assistant_node_id, "type": "llm", "data": node_data }) edge_id = str(uuid.uuid4()) self._add_graph_edge(conn, graph_id, { "id": edge_id, "source": parent_node_id, "target": assistant_node_id }) self._update_graph_current_node_id(conn, graph_id, assistant_node_id) conn.commit() return assistant_node_id except Exception as e: conn.rollback() print(f"Ошибка при создании узла-заглушки: {e}") raise def mark_node_as_placeholder_and_clear_content(self, graph_id: str, node_id: str, parent_node_id: str): """ Очищает содержимое существующего узла LLM, помечает его как заглушку, и обновляет current_node_id графа. @param graph_id: ID графа. @param node_id: ID узла, который нужно очистить и пометить как заглушку. @param parent_node_id: ID родительского узла (пользовательского), к которому будет прикреплен этот LLM-узел. """ with self._get_connection() as conn: cursor = conn.cursor() try: # 1. Получаем текущие данные узла cursor.execute("SELECT node_data FROM graph_nodes_data WHERE graph_id = ? AND node_id = ?", (graph_id, node_id)) result = cursor.fetchone() if result: node_data = json.loads(result[0]) # Очищаем контент и сбрасываем состояние node_data["message"]["content"] = "" node_data["label"] = "Получение ответа..." node_data["is_placeholder"] = True # Сбрасываем флаги генерации заголовка для перегенерации node_data["title"] = None node_data["title_generated"] = False # Обновляем узел cursor.execute("UPDATE graph_nodes_data SET node_type = ?, node_data = ?, title = NULL, title_generated = FALSE WHERE graph_id = ? AND node_id = ?", ("llm", json.dumps(node_data), graph_id, node_id)) # 2. Обновляем current_node_id графа на этот узел self._update_graph_current_node_id(conn, graph_id, node_id) conn.commit() else: raise ValueError(f"Узел с ID {node_id} не найден в графе {graph_id} для обновления.") except Exception as e: conn.rollback() print(f"Ошибка при очистке и маркировке узла {node_id} как заглушки: {e}") raise def update_assistant_node_content(self, graph_id: str, node_id: str, content: str): """ Обновляет содержимое узла ассистента после завершения стриминга. """ with self._get_connection() as conn: cursor = conn.cursor() try: cursor.execute("SELECT node_data FROM graph_nodes_data WHERE graph_id = ? AND node_id = ?", (graph_id, node_id)) result = cursor.fetchone() if result: node_data = json.loads(result[0]) node_data["message"]["content"] = content node_data["label"] = content[:30] + "..." if len(content) > 30 else content node_data["is_placeholder"] = False cursor.execute("UPDATE graph_nodes_data SET node_data = ? WHERE graph_id = ? AND node_id = ?", (json.dumps(node_data), graph_id, node_id)) conn.commit() except Exception as e: conn.rollback() print(f"Ошибка при обновлении узла: {e}") raise def mark_node_as_error(self, graph_id: str, node_id: str, error_message: str): """ Помечает узел как ошибочный. """ with self._get_connection() as conn: cursor = conn.cursor() try: cursor.execute("SELECT node_data FROM graph_nodes_data WHERE graph_id = ? AND node_id = ?", (graph_id, node_id)) result = cursor.fetchone() if result: node_data = json.loads(result[0]) node_data["message"]["content"] = error_message node_data["label"] = "Ошибка" node_data["is_placeholder"] = False cursor.execute("UPDATE graph_nodes_data SET node_type = ?, node_data = ? WHERE graph_id = ? AND node_id = ?", ("error", json.dumps(node_data), graph_id, node_id)) conn.commit() except Exception as e: conn.rollback() print(f"Ошибка при маркировке узла как ошибочного: {e}") raise def get_parent_node_id(self, graph_id: str, node_id: str) -> Optional[str]: """ Возвращает ID родительского узла. """ with self._get_connection() as conn: cursor = conn.cursor() cursor.execute("SELECT source_node_id FROM graph_edges_data WHERE graph_id = ? AND target_node_id = ?", (graph_id, node_id)) result = cursor.fetchone() return result[0] if result else None def get_node_type(self, graph_id: str, node_id: str) -> Optional[str]: """ Возвращает тип узла по его ID. @param graph_id: ID графа. @param node_id: ID узла, тип которого нужно получить. @returns: Тип узла (например, 'user', 'llm', 'error') или None, если узел не найден. """ with self._get_connection() as conn: cursor = conn.cursor() cursor.execute("SELECT node_type FROM graph_nodes_data WHERE graph_id = ? AND node_id = ?", (graph_id, node_id)) result = cursor.fetchone() return result[0] if result else None def get_graph_title_generation_status(self, graph_id: str) -> bool: """ Проверяет, был ли заголовок графа уже сгенерирован (или установлен вручную). @param graph_id: ID графа. @returns: True, если заголовок был сгенерирован/установлен, False в противном случае. """ with self._get_connection() as conn: cursor = conn.cursor() cursor.execute("SELECT title_generated FROM graphs WHERE id = ?", (graph_id,)) result = cursor.fetchone() return bool(result[0]) if result else False def update_graph_settings(self, graph_id: str, custom_system_prompt: Optional[str]) -> bool: """ Обновляет настройки графа, включая пользовательский системный промпт. Если граф не существует, он будет создан. @param graph_id: ID графа. @param custom_system_prompt: Новый пользовательский системный промпт (или None для сброса). @returns: True в случае успеха, False в противном случае. """ if not self._ensure_graph_exists(graph_id): print(f"Граф с ID {graph_id} не найден. Невозможно обновить настройки.") return False with self._get_connection() as conn: cursor = conn.cursor() try: cursor.execute( "UPDATE graphs SET custom_system_prompt = ? WHERE id = ?", (custom_system_prompt, graph_id) ) conn.commit() return cursor.rowcount > 0 except Exception as e: conn.rollback() print(f"Ошибка при обновлении системного промпта для графа {graph_id}: {e}") return False