301 lines
12 KiB
Python
301 lines
12 KiB
Python
"""Memoir service — 回忆录编排（章节生成、状态流转）；通过 MemoryService 获取 evidence。"""
|
||
|
||
import uuid
|
||
from datetime import datetime, timezone
|
||
from typing import List, Optional
|
||
|
||
from app.core.logging import get_logger
|
||
from fastapi import HTTPException
|
||
from sqlalchemy import select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
from sqlalchemy.orm import joinedload
|
||
|
||
from app.agents.prompts.memory_prompts import (
|
||
CHAPTER_CATEGORIES,
|
||
CHAPTER_ORDER,
|
||
STAGE_TO_ORDER,
|
||
)
|
||
from app.features.memoir import repo
|
||
from app.features.memoir.helpers import (
|
||
chapter_to_dict,
|
||
is_image_permanently_unavailable,
|
||
)
|
||
from app.features.memoir.models import Book, Chapter, ChapterSection, MemoirImage
|
||
from app.features.memoir.memoir_images.parser import build_initial_image_assets
|
||
from app.features.memoir.memoir_images.serializers import image_dict_to_row_kwargs
|
||
from app.features.memoir.memoir_images.prompting import MemoirImagePromptService
|
||
from app.features.memoir.memoir_images.settings import MemoirImageSettings
|
||
from app.features.memory.service import MemoryService
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
async def get_or_create_book(user_id: str, db: AsyncSession):
|
||
"""Get the user's current book or return None."""
|
||
return await repo.get_current_book(user_id, db)
|
||
|
||
|
||
class MemoirService:
|
||
def __init__(
|
||
self,
|
||
db: AsyncSession,
|
||
memory_service: Optional[MemoryService] = None,
|
||
):
|
||
self._db = db
|
||
self._memory = memory_service
|
||
|
||
async def get_evidence(self, user_id: str, query: str, *, top_k: int = 10) -> dict:
|
||
"""通过 MemoryService 获取检索证据（章节生成时优先使用）。"""
|
||
if self._memory is None:
|
||
return {
|
||
"relevant_chunks": [],
|
||
"relevant_summaries": [],
|
||
"relevant_facts": [],
|
||
"timeline_hints": [],
|
||
}
|
||
return await self._memory.retrieve(user_id, query, top_k=top_k)
|
||
|
||
async def _cleanup_unavailable_images(self, ch: Chapter) -> None:
|
||
sections = getattr(ch, "sections", None) or []
|
||
cleaned = False
|
||
for s in sections:
|
||
rec = getattr(s, "image_record", None)
|
||
if rec and is_image_permanently_unavailable(rec):
|
||
logger.info("清理不可用配图: chapter=%s, section=%s", ch.id, s.id)
|
||
s.image_id = None
|
||
await self._db.delete(rec)
|
||
cleaned = True
|
||
if cleaned:
|
||
await self._db.commit()
|
||
await self._db.refresh(ch)
|
||
|
||
async def get_current_book(self, user_id: str) -> dict:
|
||
book = await repo.get_current_book(user_id, self._db)
|
||
if not book:
|
||
return {"message": "No book found"}
|
||
return {
|
||
"id": book.id,
|
||
"title": book.title,
|
||
"total_pages": book.total_pages,
|
||
"total_words": book.total_words,
|
||
"cover_image_url": book.cover_image_url,
|
||
"has_update": book.has_update,
|
||
"last_update_chapter_id": book.last_update_chapter_id,
|
||
}
|
||
|
||
async def clear_book_update(self, user_id: str) -> dict:
|
||
book = await repo.get_current_book(user_id, self._db)
|
||
if not book:
|
||
return {"status": "ok", "message": "No book found"}
|
||
book.has_update = False
|
||
await self._db.commit()
|
||
return {"status": "ok"}
|
||
|
||
async def update_book(self, book_id: str, user_id: str, title: str) -> dict:
|
||
book = await self._db.get(Book, book_id)
|
||
if not book:
|
||
raise HTTPException(status_code=404, detail="Book not found")
|
||
if book.user_id != user_id:
|
||
raise HTTPException(status_code=403, detail="无权更新此回忆录")
|
||
book.title = title
|
||
await self._db.commit()
|
||
await self._db.refresh(book)
|
||
return {
|
||
"id": book.id,
|
||
"title": book.title,
|
||
"total_pages": book.total_pages,
|
||
"total_words": book.total_words,
|
||
"cover_image_url": book.cover_image_url,
|
||
"has_update": book.has_update,
|
||
"last_update_chapter_id": book.last_update_chapter_id,
|
||
}
|
||
|
||
async def export_pdf(self, user_id: str, book_id: str) -> dict:
|
||
from app.features.memoir.pdf_service import pdf_service
|
||
|
||
book = await self._db.get(Book, book_id)
|
||
if not book:
|
||
raise HTTPException(status_code=404, detail="Book not found")
|
||
if book.user_id != user_id:
|
||
raise HTTPException(status_code=403, detail="无权导出此回忆录")
|
||
stmt = (
|
||
select(Chapter)
|
||
.where(Chapter.user_id == user_id, Chapter.is_active == True)
|
||
.options(joinedload(Chapter.sections))
|
||
.order_by(Chapter.order_index)
|
||
)
|
||
result = await self._db.execute(stmt)
|
||
chapters = list(result.unique().scalars().all())
|
||
pdf_bytes = await pdf_service.generate_pdf(book, chapters)
|
||
return {
|
||
"pdf_base64": pdf_bytes.decode("latin1"),
|
||
"filename": f"{book.title}.pdf",
|
||
}
|
||
|
||
async def get_chapters(
|
||
self, user_id: str, is_new: bool | None = None
|
||
) -> List[dict]:
|
||
chapters = await repo.get_chapters_with_sections(
|
||
user_id, self._db, is_new_only=is_new
|
||
)
|
||
chapter_by_category: dict[str, Chapter] = {}
|
||
for ch in chapters:
|
||
if ch.category and ch.category not in chapter_by_category:
|
||
chapter_by_category[ch.category] = ch
|
||
all_chapters: List[dict] = []
|
||
for category in CHAPTER_ORDER:
|
||
ch = chapter_by_category.pop(category, None)
|
||
if ch:
|
||
await self._cleanup_unavailable_images(ch)
|
||
all_chapters.append(chapter_to_dict(ch))
|
||
else:
|
||
if is_new is True:
|
||
continue
|
||
all_chapters.append({
|
||
"id": f"placeholder_{category}",
|
||
"title": CHAPTER_CATEGORIES[category],
|
||
"content": "",
|
||
"order_index": STAGE_TO_ORDER.get(category, 999),
|
||
"status": "empty",
|
||
"category": category,
|
||
"images": [],
|
||
"cover_image": None,
|
||
"sections": [],
|
||
"updated_at": None,
|
||
"is_new": False,
|
||
"source_segments": [],
|
||
})
|
||
for ch in chapter_by_category.values():
|
||
await self._cleanup_unavailable_images(ch)
|
||
all_chapters.append(chapter_to_dict(ch))
|
||
return all_chapters
|
||
|
||
async def get_chapter(self, chapter_id: str, user_id: str) -> dict:
|
||
chapter = await repo.get_chapter_by_id(chapter_id, self._db)
|
||
if not chapter:
|
||
raise HTTPException(status_code=404, detail="Chapter not found")
|
||
if chapter.user_id != user_id:
|
||
raise HTTPException(status_code=403, detail="无权访问此章节")
|
||
if not chapter.is_active:
|
||
raise HTTPException(status_code=404, detail="Chapter not found")
|
||
await self._cleanup_unavailable_images(chapter)
|
||
return chapter_to_dict(chapter)
|
||
|
||
async def disable_chapter(self, chapter_id: str, user_id: str) -> dict:
|
||
chapter = await self._db.get(Chapter, chapter_id)
|
||
if not chapter:
|
||
raise HTTPException(status_code=404, detail="Chapter not found")
|
||
if chapter.user_id != user_id:
|
||
raise HTTPException(status_code=403, detail="无权操作此章节")
|
||
chapter.is_active = False
|
||
await self._db.commit()
|
||
return {"status": "ok", "message": "章节已清除"}
|
||
|
||
async def regenerate_chapter(self, chapter_id: str, user_id: str) -> dict:
|
||
chapter = await self._db.get(Chapter, chapter_id)
|
||
if not chapter:
|
||
raise HTTPException(status_code=404, detail="Chapter not found")
|
||
if chapter.user_id != user_id:
|
||
raise HTTPException(status_code=403, detail="无权操作此章节")
|
||
# TODO: 实现重新整理逻辑
|
||
return {"status": "ok", "message": "Chapter regeneration triggered"}
|
||
|
||
async def get_memoir_state(self, user_id: str) -> dict:
|
||
from app.features.memoir.state_service import get_or_create_state
|
||
|
||
state = await get_or_create_state(user_id, self._db)
|
||
return state.model_dump()
|
||
|
||
async def get_next_question_context(self, user_id: str) -> dict:
|
||
from app.features.memoir.state_service import get_or_create_state
|
||
|
||
state = await get_or_create_state(user_id, self._db)
|
||
return {
|
||
"current_stage": state.current_stage,
|
||
"empty_slots": state.empty_slots_for_current_stage(),
|
||
"covered_stages": state.covered_stages,
|
||
}
|
||
|
||
async def check_and_trigger_cover_generation(self, user_id: str) -> dict:
|
||
"""
|
||
检查可生成封面的章节（section 配图 > 3 且无已完成封面），
|
||
若有则触发生成任务。已有封面的章节不再检查。
|
||
"""
|
||
from app.tasks.memoir_tasks import generate_chapter_images
|
||
|
||
chapters = await repo.get_chapters_with_sections(
|
||
user_id, self._db, active_only=True, is_new_only=None
|
||
)
|
||
triggered: List[str] = []
|
||
for ch in chapters:
|
||
if not ch.category or ch.status == "empty":
|
||
continue
|
||
sections = getattr(ch, "sections", None) or []
|
||
section_image_count = sum(1 for s in sections if getattr(s, "image_id", None))
|
||
images = getattr(ch, "images", None) or []
|
||
cover_rec = next(
|
||
(m for m in images if getattr(m, "section_id", None) is None),
|
||
None,
|
||
)
|
||
if section_image_count <= 3:
|
||
continue
|
||
if cover_rec and (getattr(cover_rec, "status") or "").strip() == "completed":
|
||
continue
|
||
if cover_rec is None:
|
||
img_settings = MemoirImageSettings.from_env()
|
||
if img_settings.enabled:
|
||
now_iso = datetime.now(timezone.utc).isoformat()
|
||
cover_ph = {
|
||
"placeholder": "{{{{{{{{IMAGE:章节封面}}}}}}}}",
|
||
"description": "章节封面",
|
||
"index": 0,
|
||
}
|
||
style = MemoirImagePromptService.CATEGORY_STYLE_MAP.get(
|
||
ch.category or "", img_settings.default_style
|
||
)
|
||
cover_asset = build_initial_image_assets(
|
||
[cover_ph],
|
||
img_settings.provider,
|
||
style,
|
||
img_settings.default_size,
|
||
now_iso,
|
||
)[0]
|
||
kwargs = image_dict_to_row_kwargs(cover_asset)
|
||
cover_mi = MemoirImage(
|
||
id=str(uuid.uuid4()).replace("-", "")[:32],
|
||
chapter_id=ch.id,
|
||
section_id=None,
|
||
order_index=0,
|
||
**kwargs,
|
||
)
|
||
self._db.add(cover_mi)
|
||
await self._db.commit()
|
||
await self._db.refresh(ch)
|
||
logger.info("创建封面占位: chapter=%s", ch.id)
|
||
try:
|
||
generate_chapter_images.delay(ch.id)
|
||
triggered.append(ch.id)
|
||
logger.info("触发生成封面: chapter=%s", ch.id)
|
||
except Exception as exc:
|
||
logger.warning("封面生成任务派发失败: chapter=%s, error=%s", ch.id, exc)
|
||
return {"triggered": triggered}
|
||
|
||
async def mark_memoir_read(self, user_id: str) -> dict:
|
||
stmt = select(Chapter).where(
|
||
Chapter.user_id == user_id, Chapter.is_new == True
|
||
)
|
||
result = await self._db.execute(stmt)
|
||
for chapter in result.scalars().all():
|
||
chapter.is_new = False
|
||
stmt_book = (
|
||
select(Book)
|
||
.where(Book.user_id == user_id)
|
||
.order_by(Book.updated_at.desc())
|
||
)
|
||
result_book = await self._db.execute(stmt_book)
|
||
book = result_book.scalar_one_or_none()
|
||
if book:
|
||
book.has_update = False
|
||
await self._db.commit()
|
||
return {"status": "ok"}
|