286 lines
15 KiB
Python
286 lines
15 KiB
Python
# отключаем 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 |