diff --git a/apps/knowledge/viewsets.py b/apps/knowledge/viewsets.py index 945ad11..33914c1 100644 --- a/apps/knowledge/viewsets.py +++ b/apps/knowledge/viewsets.py @@ -2,13 +2,15 @@ from django.db.models import Q from rest_framework.exceptions import NotFound, PermissionDenied, ValidationError from rest_framework.parsers import FormParser, MultiPartParser from rest_framework.permissions import IsAuthenticated +from rest_framework.decorators import action +from rest_framework.response import Response from rest_framework.viewsets import ModelViewSet, ReadOnlyModelViewSet from apps.accounts.models import Organization, Role from apps.accounts.permissions import can_manage_organization from apps.knowledge.models import RoleRagDocument, TrainingFile from apps.knowledge.serializers import RoleRagDocumentSerializer, TrainingFileSerializer -from apps.knowledge.tasks import update_agent_prompts_from_file_task +from apps.knowledge.tasks import ingest_training_file_task, update_agent_prompts_from_file_task class TrainingFileViewSet(ModelViewSet): queryset = TrainingFile.objects.all() @@ -80,14 +82,28 @@ class TrainingFileViewSet(ModelViewSet): file_type=uploaded_file.content_type, ) + @action(detail=True, methods=['post'], url_path='retry') + def retry(self, request, *args, **kwargs): + instance: TrainingFile = self.get_object() + + if not can_manage_organization(request.user, instance.organization): + raise PermissionDenied('Permission denied') + + if instance.status != 'failed': + raise ValidationError({'status': 'Only failed files can be retried.'}) + + instance.status = 'ingesting' + instance.is_processed = False + instance.save(update_fields=['status', 'is_processed']) + ingest_training_file_task.delay(str(instance.uuid)) + + serializer = self.get_serializer(instance) + return Response(serializer.data) + def destroy(self, request, *args, **kwargs): - instance = self.get_object() + instance: TrainingFile = self.get_object() - is_uploader = instance.uploaded_by == request.user - is_org_owner = instance.organization.owner == request.user - is_org_manager = bool(request.user.is_manager) and instance.organization.members.filter(id=request.user.id).exists() - - if not (is_uploader or is_org_owner or is_org_manager): + if not can_manage_organization(request.user, instance.organization): raise PermissionDenied('Permission denied') role_uuid = str(instance.role.uuid) if instance.role_id else None diff --git a/site/src/router/api.ts b/site/src/router/api.ts index df3ac54..3d927df 100644 --- a/site/src/router/api.ts +++ b/site/src/router/api.ts @@ -122,6 +122,7 @@ export const API = { trainingFiles: { list: () => 'training-file/', byId: (uuid: string) => `training-file/${uuid}/`, + retry: (uuid: string) => `training-file/${uuid}/retry/`, }, roleRagDocuments: { list: () => 'role-rag-document/', diff --git a/site/src/views/OrganizationManage.vue b/site/src/views/OrganizationManage.vue index f7fe8d9..4822f7e 100644 --- a/site/src/views/OrganizationManage.vue +++ b/site/src/views/OrganizationManage.vue @@ -26,7 +26,7 @@ import type { User } from '../types/user' import type { InviteToken } from '../types/organization' import type { Role } from '../types/organization' import type { TrainingFile } from '../types/organization' -import { InboxOutlined, DeleteOutlined } from '@ant-design/icons-vue' +import { InboxOutlined, DeleteOutlined, ReloadOutlined } from '@ant-design/icons-vue' const route = useRoute() const router = useRouter() @@ -269,6 +269,18 @@ const deleteTrainingFile = async (uuid: string, fileName: string) => { }) } +const retryIngestion = async (uuid: string) => { + try { + const response = await apiClient.post(API.knowledge.trainingFiles.retry(uuid)) + const idx = trainingFiles.value.findIndex((f) => f.uuid === uuid) + if (idx !== -1) trainingFiles.value[idx] = response.data + message.success('Ingestion re-queued') + } catch (error) { + console.error('Failed to retry ingestion:', error) + message.error('Failed to retry ingestion') + } +} + const canDeleteTrainingFile = (record: TrainingFile): boolean => { if (auth.user?.uuid === record.uploaded_by?.uuid) return true if (organization.value?.owner?.uuid === auth.user?.uuid) return true @@ -286,8 +298,9 @@ const formatFileSize = (bytes: number) => { const trainingFileColumns = [ { title: 'File Name', - dataIndex: 'file_name', key: 'file_name', + customRender: ({ record }: { record: TrainingFile }) => + h('a', { href: record.file_url, target: '_blank', rel: 'noopener noreferrer' }, record.file_name), }, { title: 'Uploaded By', @@ -333,19 +346,35 @@ const trainingFileColumns = [ title: 'Action', key: 'action', customRender: ({ record }: { record: TrainingFile }) => { - if (canDeleteTrainingFile(record)) { - return h( - Button, - { - danger: true, - size: 'small', - icon: h(DeleteOutlined), - onClick: () => deleteTrainingFile(record.uuid, record.file_name), - }, - () => 'Delete', + const buttons: ReturnType[] = [] + if (record.status === 'failed' && canDeleteTrainingFile(record)) { + buttons.push( + h( + Button, + { + size: 'small', + icon: h(ReloadOutlined), + onClick: () => retryIngestion(record.uuid), + }, + () => 'Retry', + ), ) } - return null + if (canDeleteTrainingFile(record)) { + buttons.push( + h( + Button, + { + danger: true, + size: 'small', + icon: h(DeleteOutlined), + onClick: () => deleteTrainingFile(record.uuid, record.file_name), + }, + () => 'Delete', + ), + ) + } + return buttons.length ? h(Space, { size: 'small' }, () => buttons) : null }, }, ]