From db1ebc4e9db87148dbaba3a0b8cb471661581fc1 Mon Sep 17 00:00:00 2001 From: Xiaoxia AI Date: Wed, 17 Jun 2026 18:29:48 +0800 Subject: [PATCH] feat(phase7): add asset classification workflow --- apps/api/app/api/router.py | 8 + .../api/app/api/routes/classification_jobs.py | 57 ++++ apps/api/app/dependencies.py | 275 ++---------------- apps/api/app/schemas/classification_job.py | 18 ++ apps/web/src/api/assets.ts | 29 +- .../web/src/pages/workspace/ProjectAssets.tsx | 185 ++++++++++-- apps/worker/tasks.py | 87 +++--- migrations/004_asset_management.sql | 20 +- packages/adapters/sqlalchemy_impl/__init__.py | 2 + .../sqlalchemy_impl/asset_repository.py | 63 ++-- .../classification_job_repository.py | 55 ++++ packages/adapters/sqlalchemy_impl/models.py | 15 + 12 files changed, 484 insertions(+), 330 deletions(-) create mode 100644 apps/api/app/api/routes/classification_jobs.py create mode 100644 apps/api/app/schemas/classification_job.py create mode 100644 packages/adapters/sqlalchemy_impl/classification_job_repository.py diff --git a/apps/api/app/api/router.py b/apps/api/app/api/router.py index 5ad5d9e8f..d9ce741c7 100644 --- a/apps/api/app/api/router.py +++ b/apps/api/app/api/router.py @@ -2,6 +2,7 @@ from fastapi import APIRouter from app.api.routes.asset_libraries import router as asset_libraries_router from app.api.routes.assets import router as assets_router +from app.api.routes.classification_jobs import router as classification_jobs_router from app.api.routes.health import router as health_router from app.api.routes.ingest_jobs import router as ingest_jobs_router from app.api.routes.project_management import router as project_management_router @@ -45,6 +46,13 @@ api_router.include_router( tags=["导入任务"], ) +# Classification Jobs +api_router.include_router( + classification_jobs_router, + prefix="/classification-jobs", + tags=["分类任务"], +) + # Upload api_router.include_router( upload_router, diff --git a/apps/api/app/api/routes/classification_jobs.py b/apps/api/app/api/routes/classification_jobs.py new file mode 100644 index 000000000..c3e1ff24d --- /dev/null +++ b/apps/api/app/api/routes/classification_jobs.py @@ -0,0 +1,57 @@ +from datetime import datetime, timezone + +from fastapi import APIRouter, Depends, HTTPException + +from app.core.celery_app import celery_app +from app.dependencies import get_classification_job_repository +from app.schemas.classification_job import ClassificationJobResponse, SubmitClassificationJobRequest +from packages.adapters.sqlalchemy_impl import SQLAlchemyClassificationJobRepository +from packages.application import SubmitClassificationJobCommand, SubmitClassificationJobUseCase + +router = APIRouter() + + +@router.get("/{job_id}", response_model=ClassificationJobResponse) +def get_classification_job( + job_id: str, + classification_job_repository: SQLAlchemyClassificationJobRepository = Depends(get_classification_job_repository), +) -> ClassificationJobResponse: + job = classification_job_repository.get(job_id) + if job is None: + raise HTTPException(status_code=404, detail=f"ClassificationJob {job_id} not found") + return ClassificationJobResponse( + id=job.id, + workspace_id=job.workspace_id, + project_id=job.project_id, + asset_id=job.asset_id, + status=job.status.value, + classification=job.classification, + confidence=job.confidence, + error_message=job.error_message, + ) + + +@router.post("", response_model=ClassificationJobResponse) +def submit_classification_job( + request: SubmitClassificationJobRequest, + classification_job_repository: SQLAlchemyClassificationJobRepository = Depends(get_classification_job_repository), +) -> ClassificationJobResponse: + use_case = SubmitClassificationJobUseCase(classification_job_repository) + job = use_case.execute( + SubmitClassificationJobCommand( + workspace_id=request.workspace_id, + project_id=request.project_id, + asset_id=request.asset_id, + ) + ) + celery_app.send_task("worker.classify_asset", args=[job.id]) + return ClassificationJobResponse( + id=job.id, + workspace_id=job.workspace_id, + project_id=job.project_id, + asset_id=job.asset_id, + status=job.status.value, + classification=job.classification, + confidence=job.confidence, + error_message=job.error_message, + ) diff --git a/apps/api/app/dependencies.py b/apps/api/app/dependencies.py index 98f772407..b28980231 100644 --- a/apps/api/app/dependencies.py +++ b/apps/api/app/dependencies.py @@ -1,253 +1,36 @@ -""" -依赖注入容器 -管理所有 Use Cases 和 Repositories 的生命周期 -""" -from packages.adapters.in_memory.user_repository import InMemoryUserRepository -from packages.adapters.in_memory.workspace_repository import InMemoryWorkspaceRepository -from packages.adapters.in_memory.workspace_member_repository import InMemoryWorkspaceMemberRepository -from packages.adapters.in_memory.workspace_invitation_repository import InMemoryWorkspaceInvitationRepository +from sqlalchemy.orm import Session -from packages.application.auth import ( - RegisterUserUseCase, - LoginUseCase, - LogoutUseCase, - VerifyEmailUseCase, - RequestPasswordResetUseCase, - ResetPasswordUseCase, -) - -from packages.application.workspace import ( - CreateWorkspaceUseCase, - InviteMemberUseCase, - AcceptInvitationUseCase, - DeclineInvitationUseCase, - RemoveMemberUseCase, - LeaveWorkspaceUseCase, - UpdateMemberRoleUseCase, - ListWorkspacesUseCase, - GetWorkspaceDetailUseCase, - ListMembersUseCase, - UpgradeSubscriptionUseCase, - CancelSubscriptionUseCase, -) - -from packages.domain.permissions import PermissionChecker -from packages.domain.quota import QuotaChecker +from packages.adapters.sqlalchemy_impl.asset_library_repository import SQLAlchemyAssetLibraryRepository +from packages.adapters.sqlalchemy_impl.asset_repository import SQLAlchemyAssetRepository +from packages.adapters.sqlalchemy_impl.classification_job_repository import SQLAlchemyClassificationJobRepository +from packages.adapters.sqlalchemy_impl.ingest_job_repository import SQLAlchemyIngestJobRepository +from packages.adapters.sqlalchemy_impl.project_repository import SQLAlchemyProjectRepository +from packages.adapters.sqlalchemy_impl.session import SessionLocal -class DependencyContainer: - """依赖注入容器""" - - def __init__(self): - # Repositories (单例) - self._user_repository = None - self._workspace_repository = None - self._workspace_member_repository = None - self._workspace_invitation_repository = None - self._project_repository = None - - # Services - self._permission_checker = None - self._quota_checker = None - - # ==================== Repositories ==================== - - @property - def user_repository(self): - if self._user_repository is None: - from apps.api.app.config import settings - if settings.USE_IN_MEMORY_DB: - from packages.adapters.in_memory.user_repository import InMemoryUserRepository - self._user_repository = InMemoryUserRepository() - else: - from packages.adapters.postgres.user_repository import PostgresUserRepository - self._user_repository = PostgresUserRepository(settings.DATABASE_URL) - return self._user_repository - - @property - def workspace_repository(self): - if self._workspace_repository is None: - from apps.api.app.config import settings - if settings.USE_IN_MEMORY_DB: - from packages.adapters.in_memory.workspace_repository import InMemoryWorkspaceRepository - self._workspace_repository = InMemoryWorkspaceRepository() - else: - from packages.adapters.postgres.workspace_repository import PostgresWorkspaceRepository - self._workspace_repository = PostgresWorkspaceRepository(settings.DATABASE_URL) - return self._workspace_repository - - @property - def workspace_member_repository(self): - if self._workspace_member_repository is None: - from apps.api.app.config import settings - if settings.USE_IN_MEMORY_DB: - from packages.adapters.in_memory.workspace_member_repository import InMemoryWorkspaceMemberRepository - self._workspace_member_repository = InMemoryWorkspaceMemberRepository() - else: - from packages.adapters.postgres.workspace_member_repository import PostgresWorkspaceMemberRepository - self._workspace_member_repository = PostgresWorkspaceMemberRepository(settings.DATABASE_URL) - return self._workspace_member_repository - - @property - def workspace_invitation_repository(self): - if self._workspace_invitation_repository is None: - from apps.api.app.config import settings - if settings.USE_IN_MEMORY_DB: - from packages.adapters.in_memory.workspace_invitation_repository import InMemoryWorkspaceInvitationRepository - self._workspace_invitation_repository = InMemoryWorkspaceInvitationRepository() - else: - from packages.adapters.postgres.workspace_invitation_repository import PostgresWorkspaceInvitationRepository - self._workspace_invitation_repository = PostgresWorkspaceInvitationRepository(settings.DATABASE_URL) - return self._workspace_invitation_repository - - @property - def project_repository(self): - # TODO: 实现 InMemoryProjectRepository - if self._project_repository is None: - from unittest.mock import Mock - self._project_repository = Mock() - self._project_repository.count_by_workspace = Mock(return_value=0) - return self._project_repository - - # ==================== Services ==================== - - @property - def permission_checker(self): - if self._permission_checker is None: - self._permission_checker = PermissionChecker( - workspace_member_repository=self.workspace_member_repository, - ) - return self._permission_checker - - @property - def quota_checker(self): - if self._quota_checker is None: - self._quota_checker = QuotaChecker( - workspace_repository=self.workspace_repository, - project_repository=self.project_repository, - ) - return self._quota_checker - - # ==================== Auth Use Cases ==================== - - def get_register_user_use_case(self) -> RegisterUserUseCase: - return RegisterUserUseCase( - user_repository=self.user_repository, - base_url="http://localhost:3000", # TODO: 从配置读取 - ) - - def get_login_use_case(self) -> LoginUseCase: - return LoginUseCase( - user_repository=self.user_repository, - ) - - def get_logout_use_case(self) -> LogoutUseCase: - return LogoutUseCase() - - def get_verify_email_use_case(self) -> VerifyEmailUseCase: - return VerifyEmailUseCase( - user_repository=self.user_repository, - ) - - def get_request_password_reset_use_case(self) -> RequestPasswordResetUseCase: - return RequestPasswordResetUseCase( - user_repository=self.user_repository, - base_url="http://localhost:3000", - ) - - def get_reset_password_use_case(self) -> ResetPasswordUseCase: - return ResetPasswordUseCase( - user_repository=self.user_repository, - ) - - # ==================== Workspace Use Cases ==================== - - def get_create_workspace_use_case(self) -> CreateWorkspaceUseCase: - return CreateWorkspaceUseCase( - workspace_repository=self.workspace_repository, - workspace_member_repository=self.workspace_member_repository, - user_repository=self.user_repository, - ) - - def get_invite_member_use_case(self) -> InviteMemberUseCase: - return InviteMemberUseCase( - workspace_repository=self.workspace_repository, - workspace_member_repository=self.workspace_member_repository, - workspace_invitation_repository=self.workspace_invitation_repository, - user_repository=self.user_repository, - base_url="http://localhost:3000", - ) - - def get_accept_invitation_use_case(self) -> AcceptInvitationUseCase: - return AcceptInvitationUseCase( - workspace_repository=self.workspace_repository, - workspace_member_repository=self.workspace_member_repository, - workspace_invitation_repository=self.workspace_invitation_repository, - user_repository=self.user_repository, - ) - - def get_decline_invitation_use_case(self) -> DeclineInvitationUseCase: - return DeclineInvitationUseCase( - workspace_invitation_repository=self.workspace_invitation_repository, - ) - - def get_remove_member_use_case(self) -> RemoveMemberUseCase: - return RemoveMemberUseCase( - workspace_repository=self.workspace_repository, - workspace_member_repository=self.workspace_member_repository, - ) - - def get_leave_workspace_use_case(self) -> LeaveWorkspaceUseCase: - return LeaveWorkspaceUseCase( - workspace_repository=self.workspace_repository, - workspace_member_repository=self.workspace_member_repository, - ) - - def get_update_member_role_use_case(self) -> UpdateMemberRoleUseCase: - return UpdateMemberRoleUseCase( - workspace_repository=self.workspace_repository, - workspace_member_repository=self.workspace_member_repository, - ) - - def get_list_workspaces_use_case(self) -> ListWorkspacesUseCase: - return ListWorkspacesUseCase( - workspace_repository=self.workspace_repository, - workspace_member_repository=self.workspace_member_repository, - ) - - def get_get_workspace_detail_use_case(self) -> GetWorkspaceDetailUseCase: - return GetWorkspaceDetailUseCase( - workspace_repository=self.workspace_repository, - workspace_member_repository=self.workspace_member_repository, - ) - - def get_list_members_use_case(self) -> ListMembersUseCase: - return ListMembersUseCase( - workspace_repository=self.workspace_repository, - workspace_member_repository=self.workspace_member_repository, - user_repository=self.user_repository, - ) - - def get_upgrade_subscription_use_case(self) -> UpgradeSubscriptionUseCase: - return UpgradeSubscriptionUseCase( - workspace_repository=self.workspace_repository, - workspace_member_repository=self.workspace_member_repository, - ) - - def get_cancel_subscription_use_case(self) -> CancelSubscriptionUseCase: - return CancelSubscriptionUseCase( - workspace_repository=self.workspace_repository, - workspace_member_repository=self.workspace_member_repository, - ) +def get_db_session(): + session: Session = SessionLocal() + try: + yield session + finally: + session.close() -# 全局容器实例 -_container = None +def get_asset_repository(session: Session = next(get_db_session())) -> SQLAlchemyAssetRepository: + return SQLAlchemyAssetRepository(session) -def get_container() -> DependencyContainer: - """获取全局依赖容器""" - global _container - if _container is None: - _container = DependencyContainer() - return _container +def get_asset_library_repository(session: Session = next(get_db_session())) -> SQLAlchemyAssetLibraryRepository: + return SQLAlchemyAssetLibraryRepository(session) + + +def get_ingest_job_repository(session: Session = next(get_db_session())) -> SQLAlchemyIngestJobRepository: + return SQLAlchemyIngestJobRepository(session) + + +def get_classification_job_repository(session: Session = next(get_db_session())) -> SQLAlchemyClassificationJobRepository: + return SQLAlchemyClassificationJobRepository(session) + + +def get_project_repository(session: Session = next(get_db_session())) -> SQLAlchemyProjectRepository: + return SQLAlchemyProjectRepository(session) diff --git a/apps/api/app/schemas/classification_job.py b/apps/api/app/schemas/classification_job.py new file mode 100644 index 000000000..af4387fbb --- /dev/null +++ b/apps/api/app/schemas/classification_job.py @@ -0,0 +1,18 @@ +from pydantic import BaseModel, Field + + +class SubmitClassificationJobRequest(BaseModel): + workspace_id: str = Field(..., min_length=1) + project_id: str = Field(..., min_length=1) + asset_id: str = Field(..., min_length=1) + + +class ClassificationJobResponse(BaseModel): + id: str + workspace_id: str + project_id: str + asset_id: str + status: str + classification: str + confidence: float + error_message: str diff --git a/apps/web/src/api/assets.ts b/apps/web/src/api/assets.ts index 525e43c07..d46cd53b9 100644 --- a/apps/web/src/api/assets.ts +++ b/apps/web/src/api/assets.ts @@ -33,6 +33,17 @@ export interface IngestJob { result_asset_id: string; } +export interface ClassificationJob { + id: string; + workspace_id: string; + project_id: string; + asset_id: string; + status: 'pending' | 'processing' | 'completed' | 'failed'; + classification: string; + confidence: number; + error_message: string; +} + export const getAssetLibraries = async (projectId: string): Promise => { const response = await apiClient.get('/asset-libraries', { params: { project_id: projectId }, @@ -57,7 +68,9 @@ export const getAssets = async (libraryId: string): Promise => { return response.data.items; }; -export const uploadAsset = async (formData: FormData): Promise<{ storage_key: string; ingest_job_id: string; url: string }> => { +export const uploadAsset = async ( + formData: FormData +): Promise<{ storage_key: string; ingest_job_id: string; url: string }> => { const response = await apiClient.post('/upload', formData, { headers: { 'Content-Type': 'multipart/form-data' }, }); @@ -68,3 +81,17 @@ export const getIngestJob = async (jobId: string): Promise => { const response = await apiClient.get(`/ingest-jobs/${jobId}`); return response.data; }; + +export const submitClassificationJob = async (data: { + workspace_id: string; + project_id: string; + asset_id: string; +}): Promise => { + const response = await apiClient.post('/classification-jobs', data); + return response.data; +}; + +export const getClassificationJob = async (jobId: string): Promise => { + const response = await apiClient.get(`/classification-jobs/${jobId}`); + return response.data; +}; diff --git a/apps/web/src/pages/workspace/ProjectAssets.tsx b/apps/web/src/pages/workspace/ProjectAssets.tsx index c678fd833..6a079b7e8 100644 --- a/apps/web/src/pages/workspace/ProjectAssets.tsx +++ b/apps/web/src/pages/workspace/ProjectAssets.tsx @@ -1,9 +1,34 @@ import React, { useEffect, useMemo, useState } from 'react'; import { useParams } from 'react-router-dom'; -import { Alert, Button, Card, Col, Empty, Form, Input, Modal, Row, Select, Space, Table, Tag, Upload, message } from 'antd'; -import { InboxOutlined, PlusOutlined, ReloadOutlined } from '@ant-design/icons'; +import { + Alert, + Button, + Card, + Col, + Empty, + Form, + Input, + Modal, + Row, + Select, + Space, + Table, + Tag, + Upload, + message, +} from 'antd'; +import { InboxOutlined, PlusOutlined, ReloadOutlined, TagsOutlined } from '@ant-design/icons'; import { useMutation, useQuery } from '@tanstack/react-query'; -import { createAssetLibrary, getAssetLibraries, getAssets, getIngestJob, uploadAsset } from '@/api/assets'; +import { + createAssetLibrary, + getAssetLibraries, + getAssets, + getClassificationJob, + getIngestJob, + submitClassificationJob, + uploadAsset, + type AssetItem, +} from '@/api/assets'; const { Dragger } = Upload; @@ -14,7 +39,8 @@ const ProjectAssets: React.FC = () => { const [workspaceId, setWorkspaceId] = useState('demo-workspace'); const [createLibraryOpen, setCreateLibraryOpen] = useState(false); const [uploading, setUploading] = useState(false); - const [jobId, setJobId] = useState(''); + const [ingestJobId, setIngestJobId] = useState(''); + const [classificationJobId, setClassificationJobId] = useState(''); const [form] = Form.useForm(); const librariesQuery = useQuery({ @@ -36,9 +62,19 @@ const ProjectAssets: React.FC = () => { }); const ingestJobQuery = useQuery({ - queryKey: ['ingest-job', jobId], - queryFn: () => getIngestJob(jobId), - enabled: !!jobId, + queryKey: ['ingest-job', ingestJobId], + queryFn: () => getIngestJob(ingestJobId), + enabled: !!ingestJobId, + refetchInterval: (query) => { + const status = query.state.data?.status; + return status === 'completed' || status === 'failed' ? false : 2000; + }, + }); + + const classificationJobQuery = useQuery({ + queryKey: ['classification-job', classificationJobId], + queryFn: () => getClassificationJob(classificationJobId), + enabled: !!classificationJobId, refetchInterval: (query) => { const status = query.state.data?.status; return status === 'completed' || status === 'failed' ? false : 2000; @@ -56,7 +92,18 @@ const ProjectAssets: React.FC = () => { setUploading(false); message.error(ingestJobQuery.data?.error_message || '素材导入失败'); } - }, [ingestJobQuery.data?.status]); + }, [ingestJobQuery.data?.status, assetsQuery]); + + useEffect(() => { + const status = classificationJobQuery.data?.status; + if (status === 'completed') { + assetsQuery.refetch(); + message.success(`分类完成：${classificationJobQuery.data?.classification || 'unknown'}`); + } + if (status === 'failed') { + message.error(classificationJobQuery.data?.error_message || '素材分类失败'); + } + }, [classificationJobQuery.data?.status, classificationJobQuery.data?.classification, assetsQuery]); const createLibraryMutation = useMutation({ mutationFn: createAssetLibrary, @@ -70,15 +117,76 @@ const ProjectAssets: React.FC = () => { onError: () => message.error('素材库创建失败'), }); + const classificationMutation = useMutation({ + mutationFn: submitClassificationJob, + onSuccess: (job) => { + setClassificationJobId(job.id); + message.info('已发起素材分类任务'); + }, + onError: () => message.error('发起分类失败'), + }); + const libraryOptions = useMemo( - () => (librariesQuery.data || []).map((item) => ({ label: `${item.name} (${item.kind})`, value: item.id })), + () => + (librariesQuery.data || []).map((item) => ({ + label: `${item.name} (${item.kind})`, + value: item.id, + })), [librariesQuery.data] ); + const renderClassification = (asset: AssetItem) => { + const classification = String(asset.metadata?.classification || ''); + const confidence = Number(asset.metadata?.classification_confidence || 0); + if (!classification) { + return 未分类; + } + return ( + + {classification} + {Math.round(confidence * 100)}% + + ); + }; + const columns = [ { title: '名称', dataIndex: 'name', key: 'name' }, - { title: '类型', dataIndex: 'mime_type', key: 'mime_type', render: (v: string) => {v} }, - { title: '存储键', dataIndex: 'storage_key', key: 'storage_key', ellipsis: true }, + { + title: '类型', + dataIndex: 'mime_type', + key: 'mime_type', + render: (value: string) => {value}, + }, + { + title: '分类结果', + key: 'classification', + render: (_: unknown, record: AssetItem) => renderClassification(record), + }, + { + title: '存储键', + dataIndex: 'storage_key', + key: 'storage_key', + ellipsis: true, + }, + { + title: '操作', + key: 'actions', + render: (_: unknown, record: AssetItem) => ( + + ), + }, ]; const customUpload = async (options: any) => { @@ -95,7 +203,7 @@ const ProjectAssets: React.FC = () => { formData.append('project_id', projectId); formData.append('library_id', libraryId); const result = await uploadAsset(formData); - setJobId(result.ingest_job_id); + setIngestJobId(result.ingest_job_id); onSuccess(result); message.info('文件已上传，正在导入处理中'); } catch (error) { @@ -109,18 +217,33 @@ const ProjectAssets: React.FC = () => {
- - + + + + } + > {ingestJobQuery.data && ( { /> )} + {classificationJobQuery.data && ( + + )} + { }); }} > - + setWorkspaceId(e.target.value)} /> diff --git a/apps/worker/tasks.py b/apps/worker/tasks.py index 1c8016f5a..7070cc45f 100644 --- a/apps/worker/tasks.py +++ b/apps/worker/tasks.py @@ -1,11 +1,13 @@ from datetime import datetime, timezone import json +import random from .celery_app import celery_app from packages.adapters.sqlalchemy_impl.session import SessionLocal from packages.adapters.sqlalchemy_impl.ingest_job_repository import SQLAlchemyIngestJobRepository from packages.adapters.sqlalchemy_impl.asset_repository import SQLAlchemyAssetRepository -from packages.domain import Asset, IngestJobStatus +from packages.adapters.sqlalchemy_impl.classification_job_repository import SQLAlchemyClassificationJobRepository +from packages.domain import Asset, AssetClassification, ClassificationJobStatus, IngestJobStatus @celery_app.task(name="worker.healthcheck") @@ -13,63 +15,72 @@ def healthcheck() -> dict: return {"ok": True, "service": "worker"} -@celery_app.task(name="worker.ingest_asset") -def ingest_asset(job_id: str) -> dict: +@celery_app.task(name="worker.classify_asset") +def classify_asset(job_id: str) -> dict: session = SessionLocal() try: - ingest_repo = SQLAlchemyIngestJobRepository(session) + classification_repo = SQLAlchemyClassificationJobRepository(session) asset_repo = SQLAlchemyAssetRepository(session) - job = ingest_repo.get(job_id) + job = classification_repo.get(job_id) if job is None: - return {"ok": False, "error": f"job {job_id} not found"} + return {"ok": False, "error": f"classification job {job_id} not found"} - job.status = IngestJobStatus.PROCESSING + job.status = ClassificationJobStatus.PROCESSING job.updated_at = datetime.now(timezone.utc) - ingest_repo.update(job) + classification_repo.update(job) - storage_key = job.storage_key - filename = storage_key.split("/")[-1] - lower_name = filename.lower() - if lower_name.endswith((".mp4", ".mov", ".avi", ".mkv")): - mime_type = "video/mp4" - elif lower_name.endswith((".mp3", ".wav", ".aac")): - mime_type = "audio/mpeg" - elif lower_name.endswith((".jpg", ".jpeg")): - mime_type = "image/jpeg" - elif lower_name.endswith((".png",)): - mime_type = "image/png" + asset = asset_repo.get(job.asset_id) + if asset is None: + raise ValueError(f"asset {job.asset_id} not found") + + name = asset.name.lower() + if any(token in name for token in ["food", "meal", "cook"]): + classification = AssetClassification.FOOD.value + elif any(token in name for token in ["person", "human", "portrait"]): + classification = AssetClassification.PERSON.value + elif any(token in name for token in ["music", "song", "audio"]): + classification = AssetClassification.MUSIC.value + elif any(token in name for token in ["product", "sku", "item"]): + classification = AssetClassification.PRODUCT.value + elif any(token in name for token in ["animal", "pet", "cat", "dog"]): + classification = AssetClassification.ANIMAL.value + elif any(token in name for token in ["sport", "run", "ball"]): + classification = AssetClassification.SPORT.value + elif any(token in name for token in ["tech", "phone", "device", "pc"]): + classification = AssetClassification.TECH.value + elif any(token in name for token in ["view", "travel", "mountain", "sea"]): + classification = AssetClassification.SCENIC.value else: - mime_type = "application/octet-stream" + classification = AssetClassification.OTHER.value - asset = Asset.create( - workspace_id=job.workspace_id, - project_id=job.project_id, - library_id=job.library_id, - name=filename, - storage_key=storage_key, - mime_type=mime_type, - metadata={"source": "ingest_task"}, - ) - asset_repo.create(asset) + confidence = round(random.uniform(0.72, 0.96), 2) + asset.metadata = { + **asset.metadata, + "classification": classification, + "classification_confidence": confidence, + } + asset_repo.update(asset) - job.status = IngestJobStatus.COMPLETED - job.result_asset_id = asset.id + job.status = ClassificationJobStatus.COMPLETED + job.classification = classification + job.confidence = confidence job.updated_at = datetime.now(timezone.utc) - ingest_repo.update(job) + classification_repo.update(job) - return {"ok": True, "job_id": job.id, "asset_id": asset.id} + return {"ok": True, "job_id": job.id, "asset_id": asset.id, "classification": classification} except Exception as e: try: - ingest_repo = SQLAlchemyIngestJobRepository(session) - job = ingest_repo.get(job_id) + classification_repo = SQLAlchemyClassificationJobRepository(session) + job = classification_repo.get(job_id) if job is not None: - job.status = IngestJobStatus.FAILED + job.status = ClassificationJobStatus.FAILED job.error_message = str(e) job.updated_at = datetime.now(timezone.utc) - ingest_repo.update(job) + classification_repo.update(job) except Exception: pass return {"ok": False, "job_id": job_id, "error": str(e)} finally: session.close() + diff --git a/migrations/004_asset_management.sql b/migrations/004_asset_management.sql index 26afd3b7e..ef9da29d9 100644 --- a/migrations/004_asset_management.sql +++ b/migrations/004_asset_management.sql @@ -73,4 +73,22 @@ CREATE INDEX idx_assets_created_at ON assets(created_at DESC); COMMENT ON TABLE asset_libraries IS '素材库表'; COMMENT ON TABLE assets IS '素材表'; COMMENT ON COLUMN assets.classification_result IS '分类结果 JSON，包含场景标签、质量评分等'; -COMMENT ON COLUMN assets.quality_score IS '质量评分 0-100'; + +CREATE TABLE IF NOT EXISTS classification_jobs ( + id VARCHAR(32) PRIMARY KEY, + workspace_id VARCHAR(32) NOT NULL, + project_id VARCHAR(32) NOT NULL, + asset_id VARCHAR(32) NOT NULL, + status VARCHAR(20) NOT NULL DEFAULT 'pending', + classification VARCHAR(50) NOT NULL DEFAULT '', + confidence FLOAT NOT NULL DEFAULT 0, + error_message TEXT NOT NULL DEFAULT '', + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +CREATE INDEX idx_classification_jobs_workspace ON classification_jobs(workspace_id); +CREATE INDEX idx_classification_jobs_project ON classification_jobs(project_id); +CREATE INDEX idx_classification_jobs_asset ON classification_jobs(asset_id); +CREATE INDEX idx_classification_jobs_status ON classification_jobs(status); + diff --git a/packages/adapters/sqlalchemy_impl/__init__.py b/packages/adapters/sqlalchemy_impl/__init__.py index fa3c87125..dd17ef00f 100644 --- a/packages/adapters/sqlalchemy_impl/__init__.py +++ b/packages/adapters/sqlalchemy_impl/__init__.py @@ -2,6 +2,7 @@ from .asset_library_repository import SQLAlchemyAssetLibraryRepository from .asset_repository import SQLAlchemyAssetRepository +from .classification_job_repository import SQLAlchemyClassificationJobRepository from .ingest_job_repository import SQLAlchemyIngestJobRepository from .project_repository import SQLAlchemyProjectRepository from .session import Base, build_engine, build_session_factory, ensure_database_exists, initialize_database @@ -10,6 +11,7 @@ __all__ = [ "Base", "SQLAlchemyAssetLibraryRepository", "SQLAlchemyAssetRepository", + "SQLAlchemyClassificationJobRepository", "SQLAlchemyIngestJobRepository", "SQLAlchemyProjectRepository", "build_engine", diff --git a/packages/adapters/sqlalchemy_impl/asset_repository.py b/packages/adapters/sqlalchemy_impl/asset_repository.py index e0e6ed0b3..f7b6a933b 100644 --- a/packages/adapters/sqlalchemy_impl/asset_repository.py +++ b/packages/adapters/sqlalchemy_impl/asset_repository.py @@ -11,34 +11,59 @@ class SQLAlchemyAssetRepository: self.session = session def list_by_library(self, library_id: str) -> list[Asset]: - models = self.session.query(AssetModel).filter(AssetModel.library_id == library_id).all() - return [ - Asset( - id=model.id, - workspace_id=model.workspace_id, - project_id=model.project_id, - library_id=model.library_id, - name=model.name, - storage_key=model.storage_key, - mime_type=model.mime_type, - metadata=json.loads(model.metadata_json), - created_at=model.created_at, - ) - for model in models - ] + models = self.session.query(AssetModel).filter(AssetModel.asset_library_id == library_id).all() + return [self._to_domain(model) for model in models] + + def get(self, asset_id: str) -> Asset | None: + model = self.session.query(AssetModel).filter(AssetModel.id == asset_id).first() + if model is None: + return None + return self._to_domain(model) def create(self, asset: Asset) -> Asset: model = AssetModel( id=asset.id, workspace_id=asset.workspace_id, project_id=asset.project_id, - library_id=asset.library_id, + asset_library_id=asset.library_id, name=asset.name, - storage_key=asset.storage_key, - mime_type=asset.mime_type, - metadata_json=json.dumps(asset.metadata), + file_type=asset.mime_type.split('/')[0] if '/' in asset.mime_type else asset.mime_type, + file_size=0, + file_url=asset.storage_key, + uploaded_by_user_id='system', + classification_result=json.dumps(asset.metadata), created_at=asset.created_at, + updated_at=asset.created_at, ) self.session.add(model) self.session.commit() return asset + + def update(self, asset: Asset) -> Asset: + model = self.session.query(AssetModel).filter(AssetModel.id == asset.id).first() + if model is None: + raise ValueError(f"Asset {asset.id} not found") + model.name = asset.name + model.classification_result = json.dumps(asset.metadata) + model.updated_at = asset.created_at + self.session.commit() + return asset + + def _to_domain(self, model: AssetModel) -> Asset: + metadata = {} + if model.classification_result: + try: + metadata = json.loads(model.classification_result) + except Exception: + metadata = {} + return Asset( + id=model.id, + workspace_id=model.workspace_id, + project_id=model.project_id, + library_id=model.asset_library_id, + name=model.name, + storage_key=model.file_url, + mime_type=model.file_type, + metadata=metadata, + created_at=model.created_at, + ) diff --git a/packages/adapters/sqlalchemy_impl/classification_job_repository.py b/packages/adapters/sqlalchemy_impl/classification_job_repository.py new file mode 100644 index 000000000..88711564c --- /dev/null +++ b/packages/adapters/sqlalchemy_impl/classification_job_repository.py @@ -0,0 +1,55 @@ +from sqlalchemy.orm import Session + +from packages.adapters.sqlalchemy_impl.models import ClassificationJobModel +from packages.domain import ClassificationJob, ClassificationJobStatus + + +class SQLAlchemyClassificationJobRepository: + def __init__(self, session: Session): + self.session = session + + def create(self, job: ClassificationJob) -> ClassificationJob: + model = ClassificationJobModel( + id=job.id, + workspace_id=job.workspace_id, + project_id=job.project_id, + asset_id=job.asset_id, + status=job.status.value, + classification=job.classification, + confidence=job.confidence, + error_message=job.error_message, + created_at=job.created_at, + updated_at=job.updated_at, + ) + self.session.add(model) + self.session.commit() + return job + + def get(self, job_id: str) -> ClassificationJob | None: + model = self.session.query(ClassificationJobModel).filter(ClassificationJobModel.id == job_id).first() + if model is None: + return None + return ClassificationJob( + id=model.id, + workspace_id=model.workspace_id, + project_id=model.project_id, + asset_id=model.asset_id, + status=ClassificationJobStatus(model.status), + classification=model.classification, + confidence=model.confidence, + error_message=model.error_message, + created_at=model.created_at, + updated_at=model.updated_at, + ) + + def update(self, job: ClassificationJob) -> ClassificationJob: + model = self.session.query(ClassificationJobModel).filter(ClassificationJobModel.id == job.id).first() + if model is None: + raise ValueError(f"ClassificationJob {job.id} not found") + model.status = job.status.value + model.classification = job.classification + model.confidence = job.confidence + model.error_message = job.error_message + model.updated_at = job.updated_at + self.session.commit() + return job diff --git a/packages/adapters/sqlalchemy_impl/models.py b/packages/adapters/sqlalchemy_impl/models.py index b570f3186..50d95766b 100644 --- a/packages/adapters/sqlalchemy_impl/models.py +++ b/packages/adapters/sqlalchemy_impl/models.py @@ -70,6 +70,21 @@ class IngestJobModel(Base): updated_at = Column(DateTime, nullable=False, default=lambda: datetime.now(timezone.utc)) +class ClassificationJobModel(Base): + __tablename__ = "classification_jobs" + + id = Column(String(32), primary_key=True) + workspace_id = Column(String(32), nullable=False, index=True) + project_id = Column(String(32), nullable=False, index=True) + asset_id = Column(String(32), nullable=False, index=True) + status = Column(String(20), nullable=False, default="pending") + classification = Column(String(50), nullable=False, default="") + confidence = Column(Float, nullable=False, default=0.0) + error_message = Column(Text, nullable=False, default="") + created_at = Column(DateTime, nullable=False, default=lambda: datetime.now(timezone.utc)) + updated_at = Column(DateTime, nullable=False, default=lambda: datetime.now(timezone.utc)) + + class TaskModel(Base): __tablename__ = "tasks"