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
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":
# Если произошла ошибка во время стриминга, помечаем узел как ошибочный
if assistant_node_id:
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_id': graph_id,
'node_id': assistant_node_id,
'node_type': 'error'
})
yield f"data: {json.dumps(chunk_data)}\n\n"
return # Прекращаем генерацию при ошибке
try:
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",
"Неизвестная ошибка при стриминге"))
if graph_history_manager.socketio:
graph_history_manager.socketio.emit(
'node_type_changed', {
'graph_id': graph_id,
'node_id': assistant_node_id,
'node_type': 'error'
})
yield f"data: {json.dumps(chunk_data)}\n\n"
return
except GeneratorExit:
# НОВОЕ: Обработка прерывания генератора (клиент отключился)
print(f"⚠️ Клиент отменил стрим для узла {assistant_node_id}")
# После завершения стриминга обновляем узел полным контентом
if assistant_node_id:

View File

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