You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
60 lines
2.2 KiB
60 lines
2.2 KiB
import threading
|
|
|
|
import requests
|
|
from rest_framework import status
|
|
from rest_framework.authentication import TokenAuthentication
|
|
from rest_framework.generics import RetrieveUpdateAPIView
|
|
from rest_framework.permissions import IsAuthenticated
|
|
from rest_framework.response import Response
|
|
from rest_framework.views import APIView
|
|
from rest_framework.viewsets import ReadOnlyModelViewSet
|
|
|
|
from apps.account.permissions import IsSuperAdmin
|
|
|
|
from .models import AgentSettings, EmbeddingSession
|
|
from .serializers_admin import AgentSettingsSerializer, EmbeddingSessionSerializer
|
|
|
|
|
|
def trigger_embedding_session(session_id: int):
|
|
try:
|
|
requests.post(
|
|
"http://88.99.212.243:8098/api/sync-knowledge",
|
|
json={"session_id": session_id},
|
|
timeout=5,
|
|
)
|
|
except Exception as exc:
|
|
session = EmbeddingSession.objects.filter(pk=session_id).first()
|
|
if session and session.status == "PENDING":
|
|
session.status = "FAILED"
|
|
session.error_message = str(exc)
|
|
session.save(update_fields=["status", "error_message"])
|
|
|
|
|
|
class AdminAgentSettingsView(RetrieveUpdateAPIView):
|
|
serializer_class = AgentSettingsSerializer
|
|
permission_classes = [IsAuthenticated, IsSuperAdmin]
|
|
authentication_classes = [TokenAuthentication]
|
|
|
|
def get_object(self):
|
|
return AgentSettings.load()
|
|
|
|
|
|
class AdminEmbeddingSessionViewSet(ReadOnlyModelViewSet):
|
|
serializer_class = EmbeddingSessionSerializer
|
|
permission_classes = [IsAuthenticated, IsSuperAdmin]
|
|
authentication_classes = [TokenAuthentication]
|
|
pagination_class = None
|
|
|
|
def get_queryset(self):
|
|
return EmbeddingSession.objects.all().order_by("-created_at", "-id")
|
|
|
|
|
|
class AdminEmbeddingSessionCreateView(APIView):
|
|
permission_classes = [IsAuthenticated, IsSuperAdmin]
|
|
authentication_classes = [TokenAuthentication]
|
|
|
|
def post(self, request):
|
|
session = EmbeddingSession.objects.create()
|
|
threading.Thread(target=trigger_embedding_session, args=(session.id,), daemon=True).start()
|
|
serializer = EmbeddingSessionSerializer(session, context={"request": request})
|
|
return Response(serializer.data, status=status.HTTP_201_CREATED)
|