Add stop-generation button

This commit is contained in:
dimitrievgs 2025-10-19 21:33:05 +03:00
parent 84137bcf57
commit b0209d6b8e
2 changed files with 35 additions and 23 deletions

View File

@ -250,28 +250,32 @@ def chat_stream():
# Запускаем стриминг ответа от LLM # Запускаем стриминг ответа от LLM
accumulated_content = "" accumulated_content = ""
for chunk_data in run_agent_streaming(graph_id, user_node_id,
assistant_node_id, try:
system_prompt, model): for chunk_data in run_agent_streaming(graph_id, user_node_id,
if chunk_data.get("type") == "chunk": assistant_node_id,
accumulated_content += chunk_data.get("content", "") system_prompt, model):
yield f"data: {json.dumps(chunk_data)}\n\n" if chunk_data.get("type") == "chunk":
elif chunk_data.get("type") == "error": accumulated_content += chunk_data.get("content", "")
# Если произошла ошибка во время стриминга, помечаем узел как ошибочный yield f"data: {json.dumps(chunk_data)}\n\n"
if assistant_node_id: elif chunk_data.get("type") == "error":
graph_history_manager.mark_node_as_error( if assistant_node_id:
graph_id, assistant_node_id, graph_history_manager.mark_node_as_error(
chunk_data.get("error", graph_id, assistant_node_id,
"Неизвестная ошибка при стриминге")) chunk_data.get("error",
if graph_history_manager.socketio: "Неизвестная ошибка при стриминге"))
graph_history_manager.socketio.emit( if graph_history_manager.socketio:
'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 # Прекращаем генерацию при ошибке yield f"data: {json.dumps(chunk_data)}\n\n"
return
except GeneratorExit:
# НОВОЕ: Обработка прерывания генератора (клиент отключился)
print(f"⚠️ Клиент отменил стрим для узла {assistant_node_id}")
# После завершения стриминга обновляем узел полным контентом # После завершения стриминга обновляем узел полным контентом
if assistant_node_id: if assistant_node_id:

View File

@ -270,7 +270,7 @@ class CustomLLM:
response = self.client.chat.completions.create( response = self.client.chat.completions.create(
model=self.model_name, model=self.model_name,
messages=openai_messages, messages=openai_messages,
stream=True, # Всегда True для стриминга stream=True,
) )
for chunk in response: for chunk in response:
@ -278,6 +278,10 @@ class CustomLLM:
if chunk_message is not None: if chunk_message is not None:
yield chunk_message yield chunk_message
except GeneratorExit:
# НОВОЕ: Обработка прерывания генератора
print(f"⚠️ Стрим OpenAI прерван для модели {self.model_name}")
raise
except Exception as e: except Exception as e:
print(f"Ошибка при стриминге OpenAI API: {e}") print(f"Ошибка при стриминге OpenAI API: {e}")
raise raise
@ -301,6 +305,10 @@ class CustomLLM:
if getattr(delta, "finish_reason", None) is not None: if getattr(delta, "finish_reason", None) is not None:
break break
except GeneratorExit:
# НОВОЕ: Обработка прерывания генератора
print(f"⚠️ Стрим Mistral прерван для модели {self.model_name}")
raise
except Exception as e: except Exception as e:
print(f"Ошибка при стриминге Mistral API: {e}") print(f"Ошибка при стриминге Mistral API: {e}")
raise raise