"""资料导入候选抽取用例（异步后台任务，跨 services + ai_provider）。 触发端点在校验并置 job 为 ``extracting_*`` 后即返回；实际模型调用在后台任务中 进行（独立 session），完成后落候选并置 job 为 ``reviewing`` / ``parsed``。 仅适用单进程；多 worker 部署需替换为跨进程任务队列。 """ import asyncio from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.ai_provider import ProviderManager, ProviderManagerFactory from app.core.logging import logger from app.db.session import AsyncSessionLocal from app.models.candidates import MemoryCandidate, ProfileCandidate from app.models.imports import CharacterImportJob, ImportItem from app.models.personality import PersonalityProposal from app.modules.imports.extract import ( parse_memory_candidates, parse_personality_proposals, parse_profile_candidates, ) from app.modules.imports.prompts import ( build_memory_extraction_messages, build_personality_extraction_messages, build_profile_extraction_messages, ) from app.repositories import candidates as candidates_repo from app.repositories import imports as imports_repo from app.repositories.json_utils import dumps # 后台抽取任务强引用：asyncio.create_task 仅保留弱引用，需强引用防 GC。 _running_extraction_tasks: set[asyncio.Task[None]] = set() _provider_manager: ProviderManager | None = None def _get_provider_manager() -> ProviderManager: global _provider_manager if _provider_manager is None: _provider_manager = ProviderManagerFactory.create() return _provider_manager _PROFILE_MODEL_HINT = "personality-analysis" _MEMORY_MODEL_HINT = "memory-extraction" _PERSONALITY_MODEL_HINT = "personality-analysis" async def drain_pending_extractions() -> None: """等待所有进行中的后台抽取任务完成（测试断言 / 优雅停机用）。""" if _running_extraction_tasks: await asyncio.gather(*_running_extraction_tasks, return_exceptions=True) def start_profile_extraction(job_id: str) -> None: task = asyncio.create_task(_run_profile_extraction(job_id)) _running_extraction_tasks.add(task) task.add_done_callback(_running_extraction_tasks.discard) def start_memory_extraction(job_id: str) -> None: task = asyncio.create_task(_run_memory_extraction(job_id)) _running_extraction_tasks.add(task) task.add_done_callback(_running_extraction_tasks.discard) def start_personality_extraction(job_id: str) -> None: task = asyncio.create_task(_run_personality_extraction(job_id)) _running_extraction_tasks.add(task) task.add_done_callback(_running_extraction_tasks.discard) async def _load_job_and_items( session: AsyncSession, job_id: str ) -> tuple[CharacterImportJob | None, list[ImportItem]]: job = await imports_repo.get_job(session, job_id) if job is None: return None, [] result = await session.execute( select(ImportItem) .where(ImportItem.import_job_id == job_id) .order_by(ImportItem.sequence.asc()) ) return job, list(result.scalars().all()) async def _run_profile_extraction(job_id: str) -> None: async with AsyncSessionLocal() as session: try: job, items = await _load_job_and_items(session, job_id) if job is None: return system, user = build_profile_extraction_messages(_pairs(items)) raw = await _call_model(system, user, _PROFILE_MODEL_HINT, job.character_id) candidates = [ ProfileCandidate( character_id=job.character_id, import_job_id=job.id, field=draft.field_name, value=draft.value, confidence=draft.confidence, reason=draft.reason, status="pending", evidence_ids_json=dumps(_evidence_ids(items, draft.evidence)), ) for draft in parse_profile_candidates(raw) ] await candidates_repo.add_profile_candidates(session, candidates) job.status = "reviewing" if candidates else "parsed" await session.commit() except Exception as exc: await _fail_job(session, job_id, exc) async def _run_memory_extraction(job_id: str) -> None: async with AsyncSessionLocal() as session: try: job, items = await _load_job_and_items(session, job_id) if job is None: return system, user = build_memory_extraction_messages(_pairs(items)) raw = await _call_model(system, user, _MEMORY_MODEL_HINT, job.character_id) candidates = [ MemoryCandidate( character_id=job.character_id, import_job_id=job.id, type=draft.type, title=draft.title, content=draft.content, confidence=0.6, importance=draft.importance, sensitivity=draft.sensitivity, reason=draft.reason, status="pending", evidence_ids_json=dumps(_evidence_ids(items, draft.evidence)), ) for draft in parse_memory_candidates(raw) ] await candidates_repo.add_memory_candidates(session, candidates) job.status = "reviewing" if candidates else "parsed" await session.commit() except Exception as exc: await _fail_job(session, job_id, exc) async def _run_personality_extraction(job_id: str) -> None: async with AsyncSessionLocal() as session: try: job, items = await _load_job_and_items(session, job_id) if job is None: return system, user = build_personality_extraction_messages(_pairs(items), "") raw = await _call_model(system, user, _PERSONALITY_MODEL_HINT, job.character_id) proposals = [ PersonalityProposal( character_id=job.character_id, fields_json=dumps({draft.field_name: draft.change}), reason=draft.reason or "资料导入抽取的人格完善建议", status="pending", ) for draft in parse_personality_proposals(raw) ] session.add_all(proposals) await session.flush() job.status = "reviewing" if proposals else "parsed" await session.commit() except Exception as exc: await _fail_job(session, job_id, exc) def _pairs(items: list[ImportItem]) -> list[tuple[str, str]]: return [(item.role, item.content) for item in items] def _evidence_ids(items: list[ImportItem], indices: list[int]) -> list[str]: return [items[i].id for i in indices if 0 <= i < len(items)] async def _call_model(system: str, user: str, model_hint: str, _character_id: str) -> str: manager = _get_provider_manager() messages = manager.to_message_list([system], [("user", user)]) response = await manager.chat(messages, model_hint=model_hint) return response.content async def _fail_job(session: AsyncSession, job_id: str, exc: Exception) -> None: logger.error(f"导入候选抽取失败 job={job_id}: {exc}") try: job = await imports_repo.get_job(session, job_id) if job is not None: job.status = "failed" job.error_code = "extraction_failed" job.error_message = str(exc)[:500] await session.commit() except Exception: # noqa: BLE001 — 失败处理本身不能再抛 pass