199 lines
7.6 KiB
Python
199 lines
7.6 KiB
Python
"""资料导入候选抽取用例（异步后台任务，跨 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
|