From 45d892127dbb3309f09f72895b308c89a57c4293 Mon Sep 17 00:00:00 2001 From: dimitrievgs Date: Sat, 11 Oct 2025 23:45:33 +0300 Subject: [PATCH] work on separation / streaming --- .vscode/settings.json | 4 +- app/api.py | 40 +++++++++++++---- app/graph_history_manager.py | 8 +++- app/title_generator.py | 32 +++++++------ app/workflows.py | 87 +++++++++++++++++++++--------------- requirements.txt | 1 + 6 files changed, 114 insertions(+), 58 deletions(-) diff --git a/.vscode/settings.json b/.vscode/settings.json index ea32db3..bbdcefb 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -1,3 +1,5 @@ { - "terminal.integrated.defaultProfile.windows": "Git Bash" + "terminal.integrated.defaultProfile.windows": "Git Bash", + "python-envs.defaultEnvManager": "ms-python.python:venv", + "python-envs.pythonProjects": [] } diff --git a/app/api.py b/app/api.py index 66fa2ed..96a53f1 100644 --- a/app/api.py +++ b/app/api.py @@ -7,6 +7,7 @@ import json from flask import Flask, request, jsonify from flask_cors import CORS +from flask_socketio import SocketIO, emit from workflows import run_agent, graph_history_manager, run_agent_streaming from llm_client import MODELS # Добавляем импорт списка моделей @@ -16,6 +17,26 @@ CORS( ) # Разрешаем CORS для всех доменов (в production нужно настроить более строго) +# ----------------------------------------------- WebSocket Initialization --------------------------------------------------------------- + +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}") + +# ----------------------------------------------- API Endpoints --------------------------------------------------------------- + @api.route('/api/chat', methods=['POST']) def chat(): """API endpoint для обработки сообщений и возврата ответа.""" @@ -25,7 +46,7 @@ def chat(): parent_node_id = data.get( "parent_node_id") # Получаем parent_node_id из запроса system_prompt = data.get("system_prompt") - model = data.get("model") + model = data.get("model") if not message: return jsonify({"error": "Сообщение не может быть пустым."}, 400) @@ -140,7 +161,7 @@ def chat_stream(): 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" @@ -156,10 +177,10 @@ def chat_stream(): # После завершения стриминга обновляем узел полным контентом 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: @@ -184,13 +205,16 @@ def regenerate_message(): try: # Определяем родительский узел - 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) # Создаём новый узел-заглушку для регенерации - new_assistant_node_id = graph_history_manager.create_assistant_placeholder_node(graph_id, parent_node_id) - + 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 diff --git a/app/graph_history_manager.py b/app/graph_history_manager.py index 7e7ba52..a1d7c6a 100644 --- a/app/graph_history_manager.py +++ b/app/graph_history_manager.py @@ -29,6 +29,12 @@ class GraphHistoryManager: self.title_generator = TitleGenerator(self) self.title_generator.start() + self.socketio = None + + def set_socketio(self, socketio): + socketio = socketio + self.title_generator.socketio = socketio + def __del__(self): """ Деструктор класса. Останавливает генератор заголовков при уничтожении объекта. @@ -611,4 +617,4 @@ class GraphHistoryManager: 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 + return result[0] if result else None diff --git a/app/title_generator.py b/app/title_generator.py index 5705b37..da3328e 100644 --- a/app/title_generator.py +++ b/app/title_generator.py @@ -21,6 +21,8 @@ class TitleGenerator: self.llm = get_llm(DEFAULT_SUMMARIZATION_LLM_NAME) self.running = False self.thread: Optional[threading.Thread] = None + + self.socketio = None # Системные промпты для генерации заголовков self.graph_title_prompt = """Создай краткий заголовок (максимум 80 символов) для диалога на основе первого сообщения пользователя. @@ -185,6 +187,14 @@ class TitleGenerator: self.history_manager.update_graph_title(graph_id, title) print(f"Сгенерирован заголовок графа {graph_id}: {title}") + # ----------------------------------------------- WebSocket Notification --------------------------------------------------------------- + # Отправляем событие через WebSocket + if self.socketio: + self.socketio.emit('graph_title_updated', { + 'graph_id': graph_id, + 'title': title + }) + except Exception as e: print(f"Ошибка генерации заголовка графа {graph_id}: {e}") @@ -230,22 +240,18 @@ class TitleGenerator: self.history_manager.update_node_title(graph_id, node_id, title) print(f"Сгенерирован заголовок узла {node_id}: {title}") + # ----------------------------------------------- WebSocket Notification --------------------------------------------------------------- + # Отправляем событие через WebSocket + if self.socketio: + self.socketio.emit('node_title_updated', { + 'graph_id': graph_id, + 'node_id': node_id, + 'title': title + }) + except Exception as e: print(f"Ошибка генерации заголовка узла {node_id}: {e}") - '''def add_to_title_queue(self, - item_type: str, - graph_id: str, - node_id: str = None, - priority: int = 0): - """Добавляет элемент в очередь генерации заголовков.""" - with self._get_connection() as conn: - cursor = conn.cursor() - cursor.execute( - "INSERT INTO title_generation_queue (item_type, graph_id, node_id, priority) VALUES (?, ?, ?, ?)", - (item_type, graph_id, node_id, priority)) - conn.commit()''' - def add_node_to_queue(self, cursor, graph_id, node_id): # Добавляем в очередь заголовков с высоким приоритетом cursor.execute( diff --git a/app/workflows.py b/app/workflows.py index 594fb5e..ee79fb6 100644 --- a/app/workflows.py +++ b/app/workflows.py @@ -9,7 +9,7 @@ from langgraph.graph import StateGraph, END 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 +import uuid # Добавлено для генерации UUID from llm_client import get_llm # --- LangGraph: Построение графа --- @@ -101,15 +101,13 @@ def run_agent(user_input: str, Может продолжить существующий граф по graph_id. Возвращает обновленное состояние и ID графа. """ - + # ---------------------------------------------------- Initial State Setup ---------------------------------------------------------------- - initial_state = AgentState( - input=user_input, - parent_node_id=parent_node_id, - system_prompt=system_prompt, - selected_model=model or "gemini-2.5-flash" - ) - + initial_state = AgentState(input=user_input, + parent_node_id=parent_node_id, + system_prompt=system_prompt, + selected_model=model or "gemini-2.5-flash") + # Сохраняем исходные ID узлов и ребер для определения новых после выполнения графа original_node_ids: Set[str] = set() original_edge_ids: Set[str] = set() @@ -121,30 +119,39 @@ def run_agent(user_input: str, # Если parent_node_id передан, используем его для построения истории. loaded_graph_data = graph_history_manager.get_graph( existing_graph_id, target_node_id=parent_node_id) - + if loaded_graph_data: print(f"Продолжаю существующий граф {existing_graph_id}") # Загружаем узлы и ребра для продолжения графа - initial_state.graph_nodes = loaded_graph_data.get("graph_nodes", []) - initial_state.graph_edges = loaded_graph_data.get("graph_edges", []) - + initial_state.graph_nodes = loaded_graph_data.get( + "graph_nodes", []) + initial_state.graph_edges = loaded_graph_data.get( + "graph_edges", []) + # Сохраняем ID загруженных узлов и ребер - original_node_ids = {node['id'] for node in initial_state.graph_nodes} - original_edge_ids = {edge['id'] for edge in initial_state.graph_edges} + original_node_ids = { + node['id'] + for node in initial_state.graph_nodes + } + original_edge_ids = { + edge['id'] + for edge in initial_state.graph_edges + } # Устанавливаем parent_node_id для нового входного узла # Если parent_node_id был передан в запросе, используем его, # иначе берем current_node_id из загруженного графа. - initial_state.parent_node_id = parent_node_id or loaded_graph_data.get("current_node_id") - + initial_state.parent_node_id = parent_node_id or loaded_graph_data.get( + "current_node_id") + # История чата для текущего раунда работы LLM формируется из загруженных сообщений. - initial_state.temporary_chat_history = loaded_graph_data.get("messages", []) + initial_state.temporary_chat_history = loaded_graph_data.get( + "messages", []) # Определяем graph_id, который будет использоваться для сохранения. # Если существующий ID не передан, GraphHistoryManager сгенерирует новый. graph_id_to_save = existing_graph_id - # ---------------------------------------------------- run_agent ---------------------------------------------------------------- # Запускаем граф result = app.invoke(initial_state) @@ -152,18 +159,21 @@ def run_agent(user_input: str, # ---------------------------------------------------- Save Graph ---------------------------------------------------------------- # Определяем новые узлы и ребра для сохранения - newly_added_nodes = [node for node in final_state.graph_nodes if node['id'] not in original_node_ids] - newly_added_edges = [edge for edge in final_state.graph_edges if edge['id'] not in original_edge_ids] + newly_added_nodes = [ + node for node in final_state.graph_nodes + if node['id'] not in original_node_ids + ] + newly_added_edges = [ + edge for edge in final_state.graph_edges + if edge['id'] not in original_edge_ids + ] # Сохраняем обновленное состояние графа. # Обратите внимание: chat_history больше не сохраняется напрямую, # она является частью данных узлов в graph_nodes. final_graph_id = graph_history_manager.save_graph_changes( - graph_id_to_save, - newly_added_nodes, - newly_added_edges, - final_state.current_node_id - ) + graph_id_to_save, newly_added_nodes, newly_added_edges, + final_state.current_node_id) # ---------------------------------------------------- prepare response ---------------------------------------------------------------- # Для `messages` в `response_data` используем `final_state.temporary_chat_history`, @@ -185,27 +195,33 @@ 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): + +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) + 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"])) @@ -214,11 +230,12 @@ def run_agent_streaming(graph_id: str, user_node_id: str, assistant_node_id: str # Получаем 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) + 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 + yield {"type": "error", "error": str(e)} diff --git a/requirements.txt b/requirements.txt index 2163e8c..51a8016 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,5 +6,6 @@ langchain_google_genai flask flask_cors +flask_socketio mistralai \ No newline at end of file