From ee84617a1ce0130320bbbaf0ceaa26f0ba4a9202 Mon Sep 17 00:00:00 2001 From: dimitrievgs Date: Sat, 11 Oct 2025 00:58:41 +0300 Subject: [PATCH] work on streaming chat --- app/api.py | 99 ++++++++++++++++++++++- app/graph_history_manager.py | 149 ++++++++++++++++++++++++++++++++++- app/llm_client.py | 91 ++++++++++++++++++++- app/title_generator.py | 18 +++++ app/workflows.py | 41 +++++++++- 5 files changed, 391 insertions(+), 7 deletions(-) diff --git a/app/api.py b/app/api.py index c7107e5..66fa2ed 100644 --- a/app/api.py +++ b/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) diff --git a/app/graph_history_manager.py b/app/graph_history_manager.py index 2983a79..7e7ba52 100644 --- a/app/graph_history_manager.py +++ b/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 \ No newline at end of file diff --git a/app/llm_client.py b/app/llm_client.py index 268df65..d16f900 100644 --- a/app/llm_client.py +++ b/app/llm_client.py @@ -97,7 +97,7 @@ class CustomLLM: self.config = config self.name = config["name"] self.provider = config["provider"] - + try: if self.provider == "openai": @@ -109,7 +109,7 @@ class CustomLLM: elif self.provider == "mistralai": # Для Mistral используем прямые HTTP-запросы - self.client = Mistral(api_key=config["apiKey"], + self.client = Mistral(api_key=config["apiKey"], server_url=config["apiBase"]) self.model_name = config["model_name"] self.stream = config["stream"] @@ -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: """ diff --git a/app/title_generator.py b/app/title_generator.py index a8d2dd7..5705b37 100644 --- a/app/title_generator.py +++ b/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() \ No newline at end of file diff --git a/app/workflows.py b/app/workflows.py index 04e8547..594fb5e 100644 --- a/app/workflows.py +++ b/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() @@ -182,4 +183,42 @@ def run_agent(user_input: str, "current_node_id": final_state.current_node_id } } - return response_data \ No newline at end of file + 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)} \ No newline at end of file