736 lines
36 KiB
Python
736 lines
36 KiB
Python
# отключаем 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, Any]]: # ИЗМЕНЕНО: Dict[str, str] -> Dict[str, Any]
|
||
"""
|
||
Получает список сообщений от корневого узла до указанного узла.
|
||
Сообщения извлекаются из данных узлов, а не из отдельного поля.
|
||
"""
|
||
all_nodes = graph_data.get("graph_nodes", [])
|
||
edges = graph_data.get("graph_edges", [])
|
||
|
||
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)
|
||
|
||
messages_for_path = []
|
||
for node_id in path_to_target:
|
||
node = node_map.get(node_id)
|
||
if node and 'message' in node.get('data', {}):
|
||
original_message = node['data']['message']
|
||
|
||
message_dict = {
|
||
"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
|
||
}
|
||
|
||
# НОВОЕ: Добавляем attachments, если они есть
|
||
if "attachments" in original_message:
|
||
message_dict["attachments"] = original_message["attachments"]
|
||
|
||
messages_for_path.append(message_dict)
|
||
|
||
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,
|
||
attachments: List[Dict[str, Any]] = None) -> Dict[str, Any]: # ИЗМЕНЕНО
|
||
"""Создаёт узел пользователя с вложениями."""
|
||
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,
|
||
"attachments": attachments or [] # НОВОЕ
|
||
}
|
||
}
|
||
|
||
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
|