204 lines
9.1 KiB
Python
204 lines
9.1 KiB
Python
"""
|
||
Этот файл содержит API endpoints, созданные с использованием Flask.
|
||
Он обеспечивает взаимодействие с агентом через HTTP запросы,
|
||
включая обработку сообщений, управление графами и получение истории.
|
||
"""
|
||
|
||
import json
|
||
from flask import Flask, request, jsonify, Response, stream_with_context
|
||
from flask_cors import CORS
|
||
from flask_socketio import SocketIO
|
||
from workflows import graph_history_manager, run_agent_streaming
|
||
from llm_client import MODELS # Добавляем импорт списка моделей
|
||
|
||
api = Flask(__name__)
|
||
CORS(
|
||
api
|
||
) # Разрешаем 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/graphs', methods=['GET'])
|
||
def get_graphs():
|
||
"""API endpoint для получения списка графов."""
|
||
graphs = graph_history_manager.get_all_graphs_summary()
|
||
print("get_graphs response:", graphs) # Добавлено логирование
|
||
return jsonify(graphs)
|
||
|
||
|
||
@api.route('/api/graphs/<graph_id>', methods=['GET'])
|
||
def get_graph_by_id(graph_id):
|
||
"""API endpoint для получения данных конкретного графа по ID."""
|
||
graph_data = graph_history_manager.get_graph(graph_id)
|
||
if graph_data:
|
||
return jsonify(graph_data)
|
||
else:
|
||
return jsonify({"error": f"Граф с ID {graph_id} не найден."}, 404)
|
||
|
||
|
||
@api.route('/api/graphs/<graph_id>', methods=['DELETE'])
|
||
def delete_graph(graph_id):
|
||
"""API endpoint для удаления графа."""
|
||
if not graph_id:
|
||
return jsonify({"error": "Не указан ID графа для удаления."}, 400)
|
||
|
||
if graph_history_manager.delete_graph(graph_id):
|
||
return jsonify({"message": f"Граф {graph_id} успешно удален."})
|
||
else:
|
||
return jsonify({"error": f"Не удалось удалить граф {graph_id}."}, 500)
|
||
|
||
|
||
@api.route('/api/graphs/<graph_id>/nodes/<node_id>', methods=['DELETE'])
|
||
def delete_graph_node(graph_id, node_id):
|
||
"""API endpoint для удаления узла графа."""
|
||
if not graph_id or not node_id:
|
||
return jsonify({"error": "Не указаны ID графа или узла для удаления."}, 400)
|
||
|
||
try:
|
||
if graph_history_manager.delete_node(graph_id, node_id):
|
||
return jsonify({"message": f"Узел {node_id} успешно удален из графа {graph_id}."})
|
||
else:
|
||
return jsonify({"error": f"Не удалось удалить узел {node_id} из графа {graph_id}."}, 500)
|
||
except Exception as e:
|
||
print(f"Ошибка при удалении узла {node_id} из графа {graph_id}: {e}")
|
||
return jsonify({"error": f"Ошибка сервера при удалении узла: {str(e)}"}, 500)
|
||
|
||
|
||
@api.route('/api/messages/<graph_id>/<node_id>', methods=['GET'])
|
||
def get_messages_from_root_to_node(graph_id, node_id):
|
||
"""API endpoint для получения сообщений от корня до выбранной ноды."""
|
||
graph_data = graph_history_manager.get_graph(graph_id)
|
||
if not graph_data:
|
||
return jsonify({"error": "Граф не найден."}, 404)
|
||
|
||
target_node_id = node_id
|
||
if node_id == 'last':
|
||
target_node_id = graph_data.get('current_node_id')
|
||
if not target_node_id:
|
||
return jsonify({"error": "В графе нет активного узла, до которого можно было бы получить сообщения."}, 404)
|
||
|
||
messages = graph_history_manager.get_messages_from_root_to_node(
|
||
graph_id, graph_data, target_node_id)
|
||
return jsonify(messages)
|
||
|
||
@api.route('/api/models', methods=['GET'])
|
||
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)
|
||
|
||
|
||
@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)
|