""" Этот модуль управляет сохранением и извлечением истории графов запросов, используя SQLite базу данных для хранения данных. """ from typing import Dict, Any, List, Optional import json import sqlite3 class GraphHistoryManager: """ Менеджер для хранения истории графов запросов в SQLite. """ def __init__(self, db_path="graph_history.db"): self.db_path = db_path self._create_table(self._get_connection()) # Создаем таблицу при инициализации def _get_connection(self): """Получает соединение с базой данных.""" return sqlite3.connect(self.db_path) def _create_table(self, conn): """Создает таблицу для хранения графов, если она не существует.""" cursor = conn.cursor() cursor.execute(""" CREATE TABLE IF NOT EXISTS graphs ( id TEXT PRIMARY KEY, graph_nodes TEXT, graph_edges TEXT, current_node_id TEXT ) """) conn.commit() def _get_max_graph_id(self, conn): """Получает максимальный ID графа из базы данных.""" cursor = conn.cursor() cursor.execute( "SELECT MAX(CAST(SUBSTR(id, 7) AS INTEGER)) FROM graphs WHERE id LIKE 'graph_%'" ) result = cursor.fetchone()[0] return result if result is not None else 0 # ---------------------------------------------------- public methods ---------------------------------------------------------------- def save_graph(self, graph_data: Any) -> str: """ Сохраняет или обновляет данные графа в базе данных и возвращает ID. Не сохраняет полную историю сообщений, только структуру графа. """ conn = self._get_connection() cursor = conn.cursor() try: graph_id = graph_data.get("id") if not graph_id: max_graph_id = self._get_max_graph_id(conn) next_graph_id = max_graph_id + 1 graph_id = f"graph_{next_graph_id}" graph_data["id"] = graph_id # Преобразуем структуры данных в JSON-строки для хранения в SQLite graph_nodes_json = json.dumps(graph_data.get("graph_nodes", [])) graph_edges_json = json.dumps(graph_data.get("graph_edges", [])) current_node_id = graph_data.get("current_node_id", "") # Проверяем, существует ли уже запись с таким ID cursor.execute("SELECT id FROM graphs WHERE id = ?", (graph_id, )) existing_record = cursor.fetchone() if existing_record: # Обновляем существующую запись cursor.execute( """ UPDATE graphs SET graph_nodes = ?, graph_edges = ?, current_node_id = ? WHERE id = ? """, (graph_nodes_json, graph_edges_json, current_node_id, graph_id)) else: # Вставляем новую запись cursor.execute( """ INSERT INTO graphs (id, graph_nodes, graph_edges, current_node_id) VALUES (?, ?, ?, ?) """, (graph_id, graph_nodes_json, graph_edges_json, current_node_id)) conn.commit() print(f"Граф сохранен/обновлен: {graph_id}") return graph_id finally: conn.close() def get_graph(self, graph_id: str, target_node_id: Optional[str] = None) -> Optional[Dict[str, Any]]: """ Получает данные графа по ID из базы данных. Динамически вычисляет 'messages' до 'target_node_id'. Если target_node_id не указан, используется current_node_id из БД. """ conn = self._get_connection() cursor = conn.cursor() try: cursor.execute( "SELECT graph_nodes, graph_edges, current_node_id FROM graphs WHERE id = ?", (graph_id, )) result = cursor.fetchone() if result: graph_nodes_json, graph_edges_json, db_current_node_id = result # Преобразуем JSON-строки обратно в структуры данных Python graph_nodes = json.loads(graph_nodes_json) graph_edges = json.loads(graph_edges_json) resolved_current_node_id = target_node_id if target_node_id else db_current_node_id # Если нет узлов, значит граф пуст или некорректен if not graph_nodes: return { "id": graph_id, "messages": [], "graph_nodes": [], "graph_edges": [], "current_node_id": resolved_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 finally: conn.close() def get_all_graphs_summary(self) -> List[Dict[str, str]]: """ Возвращает краткий список всех сохраненных графов. Первое сообщение извлекается путем построения пути к current_node_id и взятия первого сообщения. """ conn = self._get_connection() cursor = conn.cursor() try: cursor.execute("SELECT id, graph_nodes, graph_edges, current_node_id FROM graphs") graphs_data = cursor.fetchall() summaries = [] for gid, nodes_json, edges_json, current_node_id in graphs_data: graph_nodes = json.loads(nodes_json) graph_edges = json.loads(edges_json) # Получаем все сообщения до 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": # Если не нашли пользовательского, берем первое ассистента first_message_content = full_messages[0].get("content", "No content") summaries.append({ "id": gid, "first_message": first_message_content }) return summaries finally: conn.close() 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' messages_for_path.append(node['data']['message']) return messages_for_path def delete_graph(self, graph_id: str) -> bool: """Удаляет граф из базы данных.""" conn = self._get_connection() cursor = conn.cursor() try: cursor.execute("DELETE FROM graphs WHERE id = ?", (graph_id, )) conn.commit() return cursor.rowcount > 0 # Возвращает True, если была удалена хотя бы одна строка except Exception as e: print(f"Ошибка при удалении графа: {e}") return False finally: conn.close()