# отключаем 40-ка строчное ограничение для этого файла # ----------------------------------------------- Imports ------------------------------------------------------------ from typing import Dict, Any, List, Optional import json import sqlite3 import uuid # Для генерации UUID # ----------------------------------------------- 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() # ----------------------------------------------- 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, timestamp DATETIME DEFAULT CURRENT_TIMESTAMP -- Добавляем метку времени для сортировки ) """) 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, -- JSON-строка для 'data' из узла 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() # _get_max_graph_id - удален, так как используем UUID для graph_id # ----------------------------------------------- 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 FROM graph_nodes_data WHERE graph_id = ?", (graph_id,)) nodes = [] for node_id, node_type, node_data_json in cursor.fetchall(): nodes.append({ "id": node_id, "type": node_type, "data": json.loads(node_data_json) }) 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 save_graph_changes(self, graph_id: str, new_nodes: List[Dict[str, Any]], new_edges: List[Dict[str, Any]], current_node_id: str) -> str: """ Сохраняет *изменения* в граф в базе данных. Добавляет новые узлы и ребра и обновляет current_node_id. Операции проводятся в рамках одной транзакции. Конфликты при добавлении существующих узлов/ребер игнорируются. Если graph_id пуст, генерируется новый UUID для графа. """ if not graph_id: graph_id = str(uuid.uuid4()) # Генерируем новый UUID для ID графа with self._get_connection() as conn: # Одна транзакция для всех изменений try: # Убедимся, что запись о графе существует в главной таблице `graphs` # INSERT OR IGNORE создаст новую запись, если её нет. # Если граф уже существует, это ничего не изменит. conn.execute("INSERT OR IGNORE INTO graphs (id) VALUES (?)", (graph_id,)) # Добавляем новые узлы for node in new_nodes: self._add_graph_node(conn, graph_id, node) # Добавляем новые ребра for edge in new_edges: self._add_graph_edge(conn, graph_id, edge) # Обновляем current_node_id - это атомарное изменение в рамках транзакции self._update_graph_current_node_id(conn, graph_id, current_node_id) conn.commit() # Фиксируем все изменения print(f"Изменения графа {graph_id} сохранены.") return graph_id except sqlite3.OperationalError as e: conn.rollback() # Откатываем транзакцию при ошибке (например, DB Locked) print(f"Ошибка блокировки SQLite при сохранении изменений графа {graph_id}: {e}") raise except Exception as e: conn.rollback() # Откатываем транзакцию при других ошибках print(f"Ошибка при сохранении изменений графа {graph_id}: {e}") raise def get_graph(self, graph_id: str, target_node_id: Optional[str] = None) -> Optional[Dict[str, Any]]: """ Получает данные графа по ID из базы данных. """ with self._get_connection() as conn: cursor = conn.cursor() cursor.execute("SELECT current_node_id FROM graphs WHERE id = ?", (graph_id,)) result = cursor.fetchone() if result: db_current_node_id = result[0] # Получаем все узлы и ребра 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_nodes": graph_nodes, "graph_edges": graph_edges}, resolved_current_node_id ) return { "id": graph_id, "messages": messages, "graph_nodes": graph_nodes, "graph_edges": graph_edges, "current_node_id": resolved_current_node_id } 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 FROM graphs ORDER BY timestamp DESC") graphs_data = cursor.fetchall() summaries = [] for gid, current_node_id in graphs_data: 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( {"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") summaries.append({ "id": gid, "first_message": first_message_content }) return summaries def get_messages_from_root_to_node( self, 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] 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"), # Используем 'type' узла "content": original_message.get("content", ""), "node_id": original_message.get("node_id", node_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