241 lines
11 KiB
Python
241 lines
11 KiB
Python
"""
|
||
Этот модуль управляет сохранением и извлечением истории графов
|
||
запросов, используя 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()
|