Add graph options with custom system prompt, make more correct handling for New Chat button

app/api.py

@@ -16,25 +16,28 @@ CORS(
     api
 )  # Разрешаем CORS для всех доменов (в production нужно настроить более строго)
 
-
 # ----------------------------------------------- WebSocket Initialization ---------------------------------------------------------------
 
-socketio = SocketIO(api, cors_allowed_origins="*")  # В production настроить CORS строже
+socketio = SocketIO(
+    api, cors_allowed_origins="*")  # В production настроить CORS строже
 
 graph_history_manager.set_socketio(socketio)
 
 # ----------------------------------------------- WebSocket Events ---------------------------------------------------------------
 
+
 @socketio.on('connect')
 def handle_connect():
     """Обработчик подключения клиента к WebSocket."""
     print(f"Client connected: {request.sid}")
 
+
 @socketio.on('disconnect')
 def handle_disconnect():
     """Обработчик отключения клиента от WebSocket."""
     print(f"Client disconnected: {request.sid}")
 
+
 @socketio.on('client_proactive_ping')
 def handle_client_proactive_ping(data=None):
     """Обработчик проактивного пинга от клиента."""
@@ -42,8 +45,10 @@ def handle_client_proactive_ping(data=None):
     # Возвращаем словарь, который будет преобразован в JSON и отправлен как ответ на callback
     return {'status': 'pong'}
 
+
 # ----------------------------------------------- API Endpoints ---------------------------------------------------------------
 
+
 @api.route('/api/graphs', methods=['GET'])
 def get_graphs():
     """API endpoint для получения списка графов."""
@@ -52,6 +57,18 @@ def get_graphs():
     return jsonify(graphs)
 
 
+@api.route('/api/graphs/new', methods=['POST'])
+def create_new_graph_endpoint():
+    """API endpoint для создания нового пустого графа."""
+    try:
+        new_graph_id = graph_history_manager.create_new_graph()
+        return jsonify({"graph_id": new_graph_id}), 201
+    except Exception as e:
+        print(f"Ошибка при создании нового графа через API: {e}")
+        return jsonify({"error": f"Не удалось создать новый граф: {str(e)}"},
+                       500)
+
+
 @api.route('/api/graphs/<graph_id>', methods=['GET'])
 def get_graph_by_id(graph_id):
     """API endpoint для получения данных конкретного графа по ID."""
@@ -87,7 +104,8 @@ def delete_all_graphs():
 def rename_graph(graph_id):
     """API endpoint для переименования графа."""
     if not graph_id:
-        return jsonify({"error": "Не указан ID графа для переименования."}, 400)
+        return jsonify({"error": "Не указан ID графа для переименования."},
+                       400)
 
     data = request.get_json()
     new_title = data.get("title")
@@ -98,23 +116,58 @@ def rename_graph(graph_id):
     if graph_history_manager.rename_graph(graph_id, new_title.strip()):
         return jsonify({"message": f"Граф {graph_id} успешно переименован."})
     else:
-        return jsonify({"error": f"Не удалось переименовать граф {graph_id}."}, 500)
+        return jsonify({"error": f"Не удалось переименовать граф {graph_id}."},
+                       500)
 
 
 @api.route('/api/graphs/<graph_id>/nodes/<node_id>', methods=['DELETE'])
 def delete_graph_node(graph_id, node_id):
     """API endpoint для удаления узла графа."""
     if not graph_id or not node_id:
-        return jsonify({"error": "Не указаны ID графа или узла для удаления."}, 400)
+        return jsonify({"error": "Не указаны ID графа или узла для удаления."},
+                       400)
 
     try:
         if graph_history_manager.delete_node(graph_id, node_id):
-            return jsonify({"message": f"Узел {node_id} успешно удален из графа {graph_id}."})
+            return jsonify({
+                "message":
+                f"Узел {node_id} успешно удален из графа {graph_id}."
+            })
         else:
-            return jsonify({"error": f"Не удалось удалить узел {node_id} из графа {graph_id}."}, 500)
+            return jsonify(
+                {
+                    "error":
+                    f"Не удалось удалить узел {node_id} из графа {graph_id}."
+                }, 500)
     except Exception as e:
         print(f"Ошибка при удалении узла {node_id} из графа {graph_id}: {e}")
-        return jsonify({"error": f"Ошибка сервера при удалении узла: {str(e)}"}, 500)
+        return jsonify(
+            {"error": f"Ошибка сервера при удалении узла: {str(e)}"}, 500)
+
+
+@api.route('/api/graphs/<graph_id>/settings', methods=['PATCH'])
+def update_graph_settings_endpoint(graph_id):
+    """API endpoint для обновления настроек графа (например, пользовательского системного промпта)."""
+    data = request.get_json()
+    custom_system_prompt = data.get(
+        "custom_system_prompt")  # Может быть None для сброса
+
+    if not graph_id:
+        return jsonify(
+            {"error": "Не указан ID графа для обновления настроек."}, 400)
+
+    try:
+        if graph_history_manager.update_graph_settings(graph_id,
+                                                       custom_system_prompt):
+            return jsonify(
+                {"message": f"Настройки графа {graph_id} успешно обновлены."})
+        else:
+            return jsonify(
+                {"error": f"Не удалось обновить настройки графа {graph_id}."},
+                500)
+    except Exception as e:
+        print(f"Ошибка при обновлении настроек графа {graph_id}: {e}")
+        return jsonify({"error": f"Ошибка сервера: {str(e)}"}, 500)
 
 
 @api.route('/api/messages/<graph_id>/<node_id>', methods=['GET'])
@@ -128,17 +181,20 @@ def get_messages_from_root_to_node(graph_id, node_id):
     if node_id == 'last':
         target_node_id = graph_data.get('current_node_id')
         if not target_node_id:
-            return jsonify({"error": "В графе нет активного узла, до которого можно было бы получить сообщения."}, 404)
+            # Если запрошен 'last', но активного узла нет, возвращаем пустой список
+            return jsonify([])
 
     messages = graph_history_manager.get_messages_from_root_to_node(
         graph_id, graph_data, target_node_id)
     return jsonify(messages)
 
+
 @api.route('/api/models', methods=['GET'])
 def get_available_models():
     """API endpoint для получения списка доступных моделей."""
     return jsonify(list(MODELS.keys()))
 
+
 @api.route('/api/chat/send', methods=['POST'])
 def send_user_message():
     """API endpoint для создания узла пользователя и получения graph_id."""
@@ -151,9 +207,11 @@ def send_user_message():
         return jsonify({"error": "Сообщение не может быть пустым."}, 400)
 
     try:
-        result = graph_history_manager.create_user_node(message, graph_id, parent_node_id)
+        result = graph_history_manager.create_user_node(
+            message, graph_id, parent_node_id)
 
-        graph_history_manager.title_generator.add_node_to_queue_direct(result.get("graph_id"), result.get("node_id"))
+        graph_history_manager.title_generator.add_node_to_queue_direct(
+            result.get("graph_id"), result.get("node_id"))
 
         return jsonify(result)
     except Exception as e:
@@ -180,42 +238,49 @@ def chat_stream():
         try:
             if existing_assistant_node_id:
                 assistant_node_id = existing_assistant_node_id
-                # НОВОЕ: Помечаем существующий узел как заглушку и очищаем его содержимое
                 graph_history_manager.mark_node_as_placeholder_and_clear_content(
-                    graph_id, assistant_node_id, user_node_id
+                    graph_id, assistant_node_id, user_node_id)
-                )
                 yield f"data: {json.dumps({'type': 'node_created', 'node_id': assistant_node_id, 'reused': True})}\n\n"
             else:
                 # Создаём узел-заглушку для ответа LLM
-                assistant_node_id = graph_history_manager.create_assistant_placeholder_node(graph_id, user_node_id)
+                assistant_node_id = graph_history_manager.create_assistant_placeholder_node(
+                    graph_id, user_node_id)
                 # Отправляем ID нового узла клиенту
                 yield f"data: {json.dumps({'type': 'node_created', 'node_id': assistant_node_id, 'reused': False})}\n\n"
 
             # Запускаем стриминг ответа от LLM
             accumulated_content = ""
-            for chunk_data in run_agent_streaming(graph_id, user_node_id, assistant_node_id, system_prompt, model):
+            for chunk_data in run_agent_streaming(graph_id, user_node_id,
+                                                  assistant_node_id,
+                                                  system_prompt, model):
                 if chunk_data.get("type") == "chunk":
                     accumulated_content += chunk_data.get("content", "")
                     yield f"data: {json.dumps(chunk_data)}\n\n"
                 elif chunk_data.get("type") == "error":
                     # Если произошла ошибка во время стриминга, помечаем узел как ошибочный
                     if assistant_node_id:
-                        graph_history_manager.mark_node_as_error(graph_id, assistant_node_id, chunk_data.get("error", "Неизвестная ошибка при стриминге"))
+                        graph_history_manager.mark_node_as_error(
+                            graph_id, assistant_node_id,
+                            chunk_data.get("error",
+                                           "Неизвестная ошибка при стриминге"))
                         if graph_history_manager.socketio:
-                            graph_history_manager.socketio.emit('node_type_changed', {
+                            graph_history_manager.socketio.emit(
-                                'graph_id': graph_id,
+                                'node_type_changed', {
-                                'node_id': assistant_node_id,
+                                    'graph_id': graph_id,
-                                'node_type': 'error'
+                                    'node_id': assistant_node_id,
-                            })
+                                    'node_type': 'error'
+                                })
                     yield f"data: {json.dumps(chunk_data)}\n\n"
-                    return # Прекращаем генерацию при ошибке
+                    return  # Прекращаем генерацию при ошибке
 
             # После завершения стриминга обновляем узел полным контентом
             if assistant_node_id:
-                graph_history_manager.update_assistant_node_content(graph_id, assistant_node_id, accumulated_content)
+                graph_history_manager.update_assistant_node_content(
+                    graph_id, assistant_node_id, accumulated_content)
 
                 # Инициируем генерацию заголовка
-                graph_history_manager.title_generator.add_node_to_queue_direct(graph_id, assistant_node_id)
+                graph_history_manager.title_generator.add_node_to_queue_direct(
+                    graph_id, assistant_node_id)
 
                 yield f"data: {json.dumps({'type': 'done', 'node_id': assistant_node_id})}\n\n"
 
@@ -223,17 +288,24 @@ def chat_stream():
             print(f"Ошибка при стриминге ответа: {e}")
             error_message = str(e)
             if assistant_node_id:
-                graph_history_manager.mark_node_as_error(graph_id, assistant_node_id, f"Ошибка при стриминге: {error_message}")
+                graph_history_manager.mark_node_as_error(
+                    graph_id, assistant_node_id,
+                    f"Ошибка при стриминге: {error_message}")
                 if graph_history_manager.socketio:
-                    graph_history_manager.socketio.emit('node_type_changed', {
+                    graph_history_manager.socketio.emit(
-                        'graph_id': graph_id,
+                        'node_type_changed', {
-                        'node_id': assistant_node_id,
+                            'graph_id': graph_id,
-                        'node_type': 'error'
+                            'node_id': assistant_node_id,
-                    })
+                            'node_type': 'error'
-            error_data = {'type': 'error', 'error': f"Ошибка при стриминге: {error_message}"}
+                        })
+            error_data = {
+                'type': 'error',
+                'error': f"Ошибка при стриминге: {error_message}"
+            }
             yield f"data: {json.dumps(error_data)}\n\n"
 
-    return Response(stream_with_context(generate()), content_type='text/event-stream')
+    return Response(stream_with_context(generate()),
+                    content_type='text/event-stream')
 
 
 @api.route('/api/chat/regenerate', methods=['POST'])
@@ -244,17 +316,24 @@ def regenerate_message():
     """
     data = request.get_json()
     graph_id = data.get("graph_id")
-    node_id = data.get("node_id") # Это ID узла, который нужно регенерировать (llm-ответ)
+    node_id = data.get(
+        "node_id")  # Это ID узла, который нужно регенерировать (llm-ответ)
     model = data.get("model")
+    system_prompt = data.get("system_prompt")
 
     if not graph_id or not node_id:
         return jsonify({"error": "Не указаны graph_id или node_id."}, 400)
 
     try:
         # Определяем родительский узел для LLM-ответа. Это должен быть user-узел.
-        parent_node_id = graph_history_manager.get_parent_node_id(graph_id, node_id)
+        parent_node_id = graph_history_manager.get_parent_node_id(
+            graph_id, node_id)
         if not parent_node_id:
-            return jsonify({"error": "Не удалось найти родительский узел для регенерации."}, 404)
+            return jsonify(
+                {
+                    "error":
+                    "Не удалось найти родительский узел для регенерации."
+                }, 404)
 
         # НОВОЕ: Проверяем тип узла
         node_type = graph_history_manager.get_node_type(graph_id, node_id)
@@ -263,16 +342,21 @@ def regenerate_message():
         if node_type == 'error':
             # Если узел был ошибочным, переиспользуем его
             # Очищаем его содержимое и помечаем как заглушку сразу же
-            assistant_node_id_to_use = node_id # Используем тот же ID
+            assistant_node_id_to_use = node_id  # Используем тот же ID
         else:
             # Если узел не был ошибочным, создаем новый узел LLM-ответа
-            assistant_node_id_to_use = graph_history_manager.create_assistant_placeholder_node(graph_id, parent_node_id)
+            assistant_node_id_to_use = graph_history_manager.create_assistant_placeholder_node(
+                graph_id, parent_node_id)
 
         return jsonify({
-            "existing_assistant_node_id": assistant_node_id_to_use, # Возвращаем ID узла, который будет регенерирован/использован
+            "existing_assistant_node_id":
-            "user_node_id": parent_node_id, # Это родительский узел (пользовательский)
+            assistant_node_id_to_use,  # Возвращаем ID узла, который будет регенерирован/использован
+            "user_node_id":
+            parent_node_id,  # Это родительский узел (пользовательский)
             "graph_id": graph_id,
-            "model": model # Передаем выбранную модель
+            "model": model,
+            "system_prompt":
+            system_prompt  # НОВОЕ: Возвращаем системный промпт, чтобы ChatView мог его передать в streamLLMResponse
         })
     except Exception as e:
         print(f"Ошибка при регенерации сообщения: {e}")

app/graph_history_manager.py

@@ -70,9 +70,19 @@ class GraphHistoryManager:
                     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,
@@ -99,7 +109,16 @@ class GraphHistoryManager:
             """)
             conn.commit()
 
-    # _get_max_graph_id - удален, так как используем UUID для graph_id
+    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]]:
@@ -157,6 +176,26 @@ class GraphHistoryManager:
 
     # ----------------------------------------------- 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,
@@ -164,14 +203,18 @@ class GraphHistoryManager:
         """
         Получает данные графа по 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 FROM graphs WHERE id = ?", # Извлекаем title и title_generated
+            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 = result # Распаковываем полученные значения
+                db_current_node_id, graph_title, graph_title_generated, custom_system_prompt = result
 
                 # Получаем все узлы и ребра
                 graph_nodes = self._get_graph_nodes(conn, graph_id)
@@ -202,8 +245,9 @@ class GraphHistoryManager:
                     "graph_nodes": graph_nodes,
                     "graph_edges": graph_edges,
                     "current_node_id": resolved_current_node_id,
-                    "title": graph_title if graph_title_generated else None, # Включаем сгенерированный заголовок
+                    "title": graph_title if graph_title_generated else None,
-                    "first_message": first_message_content # Включаем первое сообщение
+                    "first_message": first_message_content,
+                    "custom_system_prompt": custom_system_prompt
                 }
             else:
                 return None
@@ -454,12 +498,13 @@ class GraphHistoryManager:
         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:
-                conn.execute("INSERT OR IGNORE INTO graphs (id) VALUES (?)", (graph_id,))
-                
                 node_data = {
                     "label": message[:30] + "..." if len(message) > 30 else message,
                     "message": {"role": "user", "content": message, "node_id": user_node_id}
@@ -482,6 +527,9 @@ class GraphHistoryManager:
                 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,
@@ -644,3 +692,42 @@ class GraphHistoryManager:
                            (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