work on streaming chat

app/api.py

@@ -4,9 +4,10 @@
 включая обработку сообщений, управление графами и получение истории.
 """
 
+import json
 from flask import Flask, request, jsonify
 from flask_cors import CORS
-from workflows import run_agent, graph_history_manager
+from workflows import run_agent, graph_history_manager, run_agent_streaming
 from llm_client import MODELS  # Добавляем импорт списка моделей
 
 api = Flask(__name__)
@@ -101,3 +102,99 @@ def get_messages_from_root_to_node(graph_id, node_id):
 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."""
+    data = request.get_json()
+    message = data.get("message")
+    graph_id = data.get("graph_id")
+    parent_node_id = data.get("parent_node_id")
+
+    if not message:
+        return jsonify({"error": "Сообщение не может быть пустым."}, 400)
+
+    try:
+        result = graph_history_manager.create_user_node(message, graph_id, parent_node_id)
+        return jsonify(result)
+    except Exception as e:
+        print(f"Ошибка при создании узла пользователя: {e}")
+        return jsonify({"error": str(e)}, 500)
+
+
+from flask import Response, stream_with_context
+
+@api.route('/api/chat/stream', methods=['POST'])
+def chat_stream():
+    """API endpoint для стриминга ответа LLM."""
+    data = request.get_json()
+    graph_id = data.get("graph_id")
+    user_node_id = data.get("user_node_id")
+    system_prompt = data.get("system_prompt")
+    model = data.get("model")
+
+    if not graph_id or not user_node_id:
+        return jsonify({"error": "Не указаны graph_id или user_node_id."}, 400)
+
+    def generate():
+        try:
+            # Создаём узел-заглушку для ответа LLM
+            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})}\n\n"
+
+            # Запускаем стриминг ответа от LLM
+            accumulated_content = ""
+            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":
+                    yield f"data: {json.dumps(chunk_data)}\n\n"
+                    return
+
+            # После завершения стриминга обновляем узел полным контентом
+            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)
+            
+            yield f"data: {json.dumps({'type': 'done', 'node_id': assistant_node_id})}\n\n"
+
+        except Exception as e:
+            print(f"Ошибка при стриминге ответа: {e}")
+            error_data = {'type': 'error', 'error': str(e)}
+            yield f"data: {json.dumps(error_data)}\n\n"
+
+    return Response(stream_with_context(generate()), content_type='text/event-stream')
+
+
+@api.route('/api/chat/regenerate', methods=['POST'])
+def regenerate_message():
+    """API endpoint для регенерации сообщения."""
+    data = request.get_json()
+    graph_id = data.get("graph_id")
+    node_id = data.get("node_id")
+    system_prompt = data.get("system_prompt")
+    model = data.get("model")
+
+    if not graph_id or not node_id:
+        return jsonify({"error": "Не указаны graph_id или node_id."}, 400)
+
+    try:
+        # Определяем родительский узел
+        parent_node_id = graph_history_manager.get_parent_node_id(graph_id, node_id)
+        if not parent_node_id:
+            return jsonify({"error": "Не удалось найти родительский узел."}, 404)
+
+        # Создаём новый узел-заглушку для регенерации
+        new_assistant_node_id = graph_history_manager.create_assistant_placeholder_node(graph_id, parent_node_id)
+        
+        return jsonify({
+            "new_node_id": new_assistant_node_id,
+            "parent_node_id": parent_node_id
+        })
+    except Exception as e:
+        print(f"Ошибка при регенерации сообщения: {e}")
+        return jsonify({"error": str(e)}, 500)

app/graph_history_manager.py

@@ -425,10 +425,11 @@ class GraphHistoryManager:
                     new_current_node_id = None
                     if parents:
                         # Если есть родители, делаем текущим узлом одного из родителей
-                        # Для простоты, берем первого родителя. Можно добавить более сложную логику,
-                        # например, последнего добавленного родителя, если timestamp доступен.
                         new_current_node_id = parents[0]
-                    # Если детей нет и родителей тоже, current_node_id останется None.
+                    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))
@@ -469,3 +470,145 @@ class GraphHistoryManager:
             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) -> Dict[str, Any]:
+        """
+        Создаёт узел пользователя и возвращает graph_id и node_id.
+        """
+        if not graph_id:
+            graph_id = str(uuid.uuid4())
+        
+        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}
+                }
+                
+                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()
+                
+                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 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

app/llm_client.py

@@ -239,6 +239,93 @@ class CustomLLM:
                 raise
         return ""
 
+    def stream2(self, messages: List[Any]):
+        """
+        Генератор для стриминга ответа от LLM.
+        Возвращает чанки контента по мере их получения.
+        """
+        if self.provider == "openai":
+            openai_messages = self._prepare_openai_messages(messages)
+            
+            try:
+                response = self.client.chat.completions.create(
+                    model=self.model_name,
+                    messages=openai_messages,
+                    stream=True,  # Всегда True для стриминга
+                )
+
+                for chunk in response:
+                    chunk_message = chunk.choices[0].delta.content
+                    if chunk_message is not None:
+                        yield chunk_message
+
+            except Exception as e:
+                print(f"Ошибка при стриминге OpenAI API: {e}")
+                raise
+
+        elif self.provider == "mistralai":
+            mistralai_messages = self._prepare_mistralai_messages(messages)
+            
+            try:
+                response = self.client.chat.stream(
+                    model=self.model_name,
+                    messages=mistralai_messages
+                )
+
+                for chunk in response:
+                    delta = chunk.data.choices[0].delta
+                    chunk_message = getattr(delta, "content", None)
+                    
+                    if chunk_message is not None:
+                        yield chunk_message
+                    
+                    if getattr(delta, "finish_reason", None) is not None:
+                        break
+
+            except Exception as e:
+                print(f"Ошибка при стриминге Mistral API: {e}")
+                raise
+
+    def _prepare_openai_messages(self, messages: List[Any]) -> List[Dict[str, str]]:
+        """Подготавливает сообщения в формате OpenAI."""
+        openai_messages = []
+        if isinstance(messages, str):
+            openai_messages.append({"role": "user", "content": messages})
+        elif isinstance(messages, list):
+            for msg in messages:
+                if isinstance(msg, HumanMessage):
+                    openai_messages.append({"role": "user", "content": msg.content})
+                elif isinstance(msg, SystemMessage):
+                    openai_messages.append({"role": "system", "content": msg.content})
+                elif isinstance(msg, AIMessage):
+                    openai_messages.append({"role": "assistant", "content": msg.content})
+                else:
+                    openai_messages.append({
+                        "role": msg.type if hasattr(msg, 'type') else 'user',
+                        "content": msg.content
+                    })
+        return openai_messages
+
+    def _prepare_mistralai_messages(self, messages: List[Any]) -> List[Dict[str, str]]:
+        """Подготавливает сообщения в формате Mistral AI."""
+        mistralai_messages = []
+        if isinstance(messages, str):
+            mistralai_messages.append({"role": "user", "content": messages})
+        elif isinstance(messages, list):
+            for msg in messages:
+                if isinstance(msg, HumanMessage):
+                    mistralai_messages.append({"role": "user", "content": msg.content})
+                elif isinstance(msg, SystemMessage):
+                    mistralai_messages.append({"role": "system", "content": msg.content})
+                elif isinstance(msg, AIMessage):
+                    mistralai_messages.append({"role": "assistant", "content": msg.content})
+                else:
+                    mistralai_messages.append({
+                        "role": msg.type if hasattr(msg, 'type') else 'user',
+                        "content": msg.content
+                    })
+        return mistralai_messages
+
 
 def get_llm(name: str) -> CustomLLM:
     """

app/title_generator.py

@@ -256,3 +256,21 @@ class TitleGenerator:
         cursor.execute(
             "INSERT INTO title_generation_queue (item_type, graph_id, priority) VALUES (?, ?, ?)",
             ("graph", graph_id, 10))
+
+    def add_node_to_queue_direct(self, graph_id: str, node_id: str, priority: int = 10):
+        """Добавляет узел в очередь генерации заголовков напрямую (с собственным соединением)."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                "INSERT INTO title_generation_queue (item_type, graph_id, node_id, priority) VALUES (?, ?, ?, ?)",
+                ("node", graph_id, node_id, priority))
+            conn.commit()
+
+    def add_graph_to_queue_direct(self, graph_id: str, priority: int = 10):
+        """Добавляет граф в очередь генерации заголовков напрямую (с собственным соединением)."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                "INSERT INTO title_generation_queue (item_type, graph_id, priority) VALUES (?, ?, ?)",
+                ("graph", graph_id, priority))
+            conn.commit()

app/workflows.py

@@ -10,6 +10,7 @@ from graph_history_manager import GraphHistoryManager
 from models import AgentState
 from nodes import parse_command_node, execute_command_node, call_llm_node, generate_images_node, analyze_image_node, get_meet_subtitles_node, get_teams_subtitles_node, summarize_history_node, handle_error_node, help_node
 import uuid # Добавлено для генерации UUID
+from llm_client import get_llm
 # --- LangGraph: Построение графа ---
 
 graph_history_manager = GraphHistoryManager()
@@ -183,3 +184,41 @@ def run_agent(user_input: str,
         }
     }
     return response_data
+
+def run_agent_streaming(graph_id: str, user_node_id: str, assistant_node_id: str, system_prompt: Optional[str] = None, model: Optional[str] = None):
+    """
+    Генератор для стримингового ответа LLM.
+    Возвращает чанки контента по мере их получения.
+    """
+    try:
+        # Загружаем граф и историю до user_node_id
+        loaded_graph_data = graph_history_manager.get_graph(graph_id, target_node_id=user_node_id)
+        if not loaded_graph_data:
+            yield {"type": "error", "error": "Граф не найден."}
+            return
+
+        messages = loaded_graph_data.get("messages", [])
+        
+        # Формируем сообщения для LLM
+        from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
+        
+        messages_for_llm = []
+        if system_prompt:
+            messages_for_llm.append(SystemMessage(content=system_prompt))
+        
+        for msg in messages:
+            if msg["role"] == "user":
+                messages_for_llm.append(HumanMessage(content=msg["content"]))
+            elif msg["role"] == "assistant":
+                messages_for_llm.append(AIMessage(content=msg["content"]))
+
+        # Получаем LLM и стримим ответ
+        llm = get_llm(model or "gemini-2.5-flash")
+        
+        for chunk in llm.stream2(messages_for_llm):
+            content = chunk.content if hasattr(chunk, 'content') else str(chunk)
+            yield {"type": "chunk", "content": content}
+            
+    except Exception as e:
+        print(f"Ошибка при стриминге: {e}")
+        yield {"type": "error", "error": str(e)}