初始化提交

2026-08-21 16:00:13 +08:00
commit 887c1f449e
@@ -0,0 +1,198 @@
+"""资料导入候选抽取用例（异步后台任务，跨 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