- Judge baseline excerpt and library chapter separately; build_memoir_compare_summary for gate, nine-dim and leaf deltas. - Memoir SSE chapter payload: baseline_judge, compare_summary, baseline_judge_error. - MemoirJudgeOutput: loose score coercion and post-validate clamp; memoir judge prompt caps from settings. - app-eval-web: two-column MemoirScoreCard layout, MemoirCompareSummary, chapter blocks and CSS. - Add memoir_compare_summary, log_events, celery_log_context, memoir_pipeline_progress; tests and migration 0014. - Misc: memory/evidence and enrichment paths, task/orchestrator updates, internal-eval docs, env examples.
334 lines
12 KiB
Python
334 lines
12 KiB
Python
"""
|
||
MemoirOrchestrator:按 segment 编排流水线,调用各 Specialist Agent。
|
||
负责:遍历 segments、按 category 聚合、调用 Specialist、更新 state;
|
||
持久化与章节生成由 process_category 回调完成。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import time
|
||
from dataclasses import dataclass
|
||
from typing import Any, Callable, Dict, List, Optional, Set, Tuple
|
||
|
||
from app.agents.memoir.batch_phase1_prep import (
|
||
STAGE_ALLOWED_SLOTS,
|
||
run_batch_phase1_prep_chunked,
|
||
)
|
||
from app.agents.memoir.classification_agent import (
|
||
ClassificationAgent,
|
||
_looks_like_fragment_only,
|
||
)
|
||
from app.agents.memoir.classification_agent import (
|
||
_detect_stage as detect_stage_from_keywords,
|
||
)
|
||
from app.agents.memoir.extraction_agent import ExtractionAgent, ExtractionResult
|
||
from app.agents.stage_constants import normalize_chapter_category, normalize_chat_stage
|
||
from app.agents.state_schema import MemoirStateSchema
|
||
from app.core.agent_logging import agent_span, agent_summary_enabled, log_agent_detail
|
||
from app.core.config import settings
|
||
from app.core.logging import get_logger
|
||
from app.features.conversation.models import Segment
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
@dataclass
|
||
class PreparedMemoirBatches:
|
||
"""Explicit batching result: updated state + segments grouped by chapter category."""
|
||
|
||
state: MemoirStateSchema
|
||
category_to_segments: Dict[str, List[Segment]]
|
||
#: segment id 在「LLM 判 none 且 extraction slots 为空」时加入;batch 级短路见 memoir_tasks
|
||
segment_skip_story_ids: Set[str]
|
||
#: 每个 segment → Phase 1 分类 chapter_category(持久化到 Segment.topic_category)
|
||
segment_chapter_category: Dict[str, str]
|
||
|
||
|
||
class MemoirOrchestrator:
|
||
"""
|
||
回忆录生成编排器。
|
||
遍历 segments → ExtractionAgent → ClassificationAgent → 按 category 聚合 →
|
||
调用 process_category 生成叙事并持久化。
|
||
"""
|
||
|
||
def __init__(self) -> None:
|
||
self.extraction_agent = ExtractionAgent()
|
||
self.classification_agent = ClassificationAgent()
|
||
|
||
def prepare_batches(
|
||
self,
|
||
*,
|
||
segments: List[Segment],
|
||
llm: Any,
|
||
get_or_create_state: Callable[[], MemoirStateSchema],
|
||
update_slot: Callable[[str, str, str, List[str]], MemoirStateSchema],
|
||
llm_fast: Any | None = None,
|
||
on_phase1_chunk: Optional[Callable[[int, int], None]] = None,
|
||
) -> PreparedMemoirBatches:
|
||
"""
|
||
遍历 segments:Extraction → slot 更新 → Classification → 按 category 分桶。
|
||
不含锁与写章节/故事(由调用方显式执行)。
|
||
|
||
``llm_fast``:分类与抽取专用;未传时与 ``llm`` 相同(叙事/路由仍用 ``llm``)。
|
||
"""
|
||
state = get_or_create_state()
|
||
category_to_segments: Dict[str, List[Segment]] = {}
|
||
segment_skip_story_ids: Set[str] = set()
|
||
segment_chapter_category: Dict[str, str] = {}
|
||
classify_extract_llm = llm_fast if llm_fast is not None else llm
|
||
|
||
# batch 路径为默认主路径(需 LLM + 开关),失败自动回退逐段
|
||
use_batch = (
|
||
bool(segments)
|
||
and classify_extract_llm is not None
|
||
and settings.memoir_phase1_batch_llm_enabled
|
||
)
|
||
if use_batch:
|
||
try:
|
||
result = self._prepare_batches_via_batch_llm(
|
||
segments=segments,
|
||
state=state,
|
||
classify_extract_llm=classify_extract_llm,
|
||
update_slot=update_slot,
|
||
on_phase1_chunk=on_phase1_chunk,
|
||
)
|
||
logger.info(
|
||
"event=phase1_batch_path_used segment_count={} "
|
||
"msg=Phase1 批处理 LLM 路径已使用",
|
||
len(segments),
|
||
)
|
||
return result
|
||
except Exception as e:
|
||
logger.warning(
|
||
"event=phase1_batch_path_fallback segment_count={} exc={} "
|
||
"msg=Phase1 批处理失败,回退逐段",
|
||
len(segments),
|
||
e,
|
||
)
|
||
|
||
for segment in segments:
|
||
text = segment.user_input_text or ""
|
||
seg_t0 = time.perf_counter()
|
||
initial_stage = detect_stage_from_keywords(
|
||
text, state.current_stage or "childhood"
|
||
)
|
||
stage_slots_raw = state.slots.get(initial_stage, {}) or {}
|
||
|
||
with agent_span(
|
||
logger,
|
||
"MemoirOrchestrator.ExtractionAgent.extract",
|
||
segment_id=segment.id,
|
||
):
|
||
result: ExtractionResult = self.extraction_agent.extract(
|
||
user_message=text,
|
||
current_stage=state.current_stage or "childhood",
|
||
stage_slots=stage_slots_raw,
|
||
llm=classify_extract_llm,
|
||
)
|
||
detected_stage = result.detected_stage
|
||
for slot_name, snippet in result.slots.items():
|
||
state = update_slot(detected_stage, slot_name, snippet, [segment.id])
|
||
|
||
with agent_span(
|
||
logger,
|
||
"MemoirOrchestrator.ClassificationAgent.classify",
|
||
segment_id=segment.id,
|
||
):
|
||
classify_result = self.classification_agent.classify(
|
||
text=text,
|
||
fallback_stage=detected_stage,
|
||
llm=classify_extract_llm,
|
||
segment_id=segment.id,
|
||
)
|
||
chapter_category = classify_result.category
|
||
if (not result.slots) and classify_result.llm_said_none:
|
||
segment_skip_story_ids.add(str(segment.id))
|
||
segment_chapter_category[str(segment.id)] = chapter_category
|
||
|
||
if agent_summary_enabled():
|
||
logger.info(
|
||
"MemoirOrchestrator.segment segment_id={} text_len={} "
|
||
"detected_stage={} category={} segment_total_ms={:.2f}",
|
||
segment.id,
|
||
len(text),
|
||
detected_stage,
|
||
chapter_category,
|
||
(time.perf_counter() - seg_t0) * 1000,
|
||
)
|
||
log_agent_detail(
|
||
logger,
|
||
"MemoirOrchestrator.segment_done segment_id={} slots={}",
|
||
segment.id,
|
||
list((result.slots or {}).keys()),
|
||
)
|
||
category_to_segments.setdefault(chapter_category, []).append(segment)
|
||
|
||
return PreparedMemoirBatches(
|
||
state=state,
|
||
category_to_segments=category_to_segments,
|
||
segment_skip_story_ids=segment_skip_story_ids,
|
||
segment_chapter_category=segment_chapter_category,
|
||
)
|
||
|
||
def _prepare_batches_via_batch_llm(
|
||
self,
|
||
*,
|
||
segments: List[Segment],
|
||
state: MemoirStateSchema,
|
||
classify_extract_llm: Any,
|
||
update_slot: Callable[[str, str, str, List[str]], MemoirStateSchema],
|
||
on_phase1_chunk: Optional[Callable[[int, int], None]] = None,
|
||
) -> PreparedMemoirBatches:
|
||
category_to_segments: Dict[str, List[Segment]] = {}
|
||
segment_skip_story_ids: Set[str] = set()
|
||
segment_chapter_category: Dict[str, str] = {}
|
||
|
||
by_id = run_batch_phase1_prep_chunked(
|
||
segments,
|
||
state,
|
||
classify_extract_llm,
|
||
chunk_size=int(settings.memoir_phase1_batch_llm_chunk_size),
|
||
on_chunk=on_phase1_chunk,
|
||
)
|
||
|
||
for segment in segments:
|
||
text = segment.user_input_text or ""
|
||
seg_t0 = time.perf_counter()
|
||
row = by_id[str(segment.id)]
|
||
|
||
result_slots = dict(row.slots)
|
||
fb = state.current_stage or "childhood"
|
||
|
||
if not result_slots:
|
||
detected_stage = normalize_chat_stage(fb, fb)
|
||
else:
|
||
detected_stage = normalize_chat_stage(row.detected_stage, fb)
|
||
|
||
allowed = STAGE_ALLOWED_SLOTS.get(detected_stage, frozenset())
|
||
result_slots = {k: v for k, v in result_slots.items() if k in allowed}
|
||
if not result_slots:
|
||
detected_stage = normalize_chat_stage(fb, fb)
|
||
|
||
with agent_span(
|
||
logger,
|
||
"MemoirOrchestrator.BatchPhase1Prep.apply",
|
||
segment_id=segment.id,
|
||
):
|
||
for slot_name, snippet in result_slots.items():
|
||
state = update_slot(
|
||
detected_stage, slot_name, snippet, [segment.id]
|
||
)
|
||
|
||
if _looks_like_fragment_only(text):
|
||
chapter_category = "summary"
|
||
llm_said_none = False
|
||
else:
|
||
raw_cat = (row.chapter_category_raw or "").strip().lower()
|
||
if raw_cat == "none":
|
||
chapter_category = "summary"
|
||
llm_said_none = True
|
||
else:
|
||
chapter_category = normalize_chapter_category(
|
||
row.chapter_category_raw,
|
||
"summary",
|
||
)
|
||
llm_said_none = False
|
||
|
||
if (not result_slots) and llm_said_none:
|
||
segment_skip_story_ids.add(str(segment.id))
|
||
segment_chapter_category[str(segment.id)] = chapter_category
|
||
|
||
if agent_summary_enabled():
|
||
logger.info(
|
||
"MemoirOrchestrator.segment(batch) segment_id={} text_len={} "
|
||
"detected_stage={} category={} segment_total_ms={:.2f}",
|
||
segment.id,
|
||
len(text),
|
||
detected_stage,
|
||
chapter_category,
|
||
(time.perf_counter() - seg_t0) * 1000,
|
||
)
|
||
log_agent_detail(
|
||
logger,
|
||
"MemoirOrchestrator.segment_done(batch) segment_id={} slots={}",
|
||
segment.id,
|
||
list(result_slots.keys()),
|
||
)
|
||
category_to_segments.setdefault(chapter_category, []).append(segment)
|
||
|
||
return PreparedMemoirBatches(
|
||
state=state,
|
||
category_to_segments=category_to_segments,
|
||
segment_skip_story_ids=segment_skip_story_ids,
|
||
segment_chapter_category=segment_chapter_category,
|
||
)
|
||
|
||
def run(
|
||
self,
|
||
*,
|
||
segments: List[Segment],
|
||
llm: Any,
|
||
user_profile: str = "",
|
||
user_birth_year: Any = None,
|
||
get_or_create_state: Callable[[], MemoirStateSchema],
|
||
update_slot: Callable[[str, str, str, List[str]], MemoirStateSchema],
|
||
acquire_lock: Callable[[str], bool],
|
||
release_lock: Callable[[str], None],
|
||
process_category: Callable[
|
||
[
|
||
str,
|
||
List[Segment],
|
||
MemoirStateSchema,
|
||
str,
|
||
Any,
|
||
Any,
|
||
],
|
||
Tuple[Any, bool],
|
||
],
|
||
raise_retry: Callable[[], None],
|
||
llm_fast: Any | None = None,
|
||
) -> Tuple[Set[str], int]:
|
||
"""
|
||
执行回忆录流水线。
|
||
process_category(category, segments, state, user_profile, user_birth_year, llm)
|
||
返回 (chapter, has_images_to_generate)。
|
||
返回 (chapters_to_enqueue, processed_count)。
|
||
raise_retry 用于锁竞争时抛出 Celery retry。
|
||
"""
|
||
prepared = self.prepare_batches(
|
||
segments=segments,
|
||
llm=llm,
|
||
llm_fast=llm_fast,
|
||
get_or_create_state=get_or_create_state,
|
||
update_slot=update_slot,
|
||
on_phase1_chunk=None,
|
||
)
|
||
state = prepared.state
|
||
chapters_to_enqueue: Set[str] = set()
|
||
category_to_segments = prepared.category_to_segments
|
||
|
||
# 按 category 调用 process_category:叙事生成、持久化、封面入队标记
|
||
for chapter_category, category_segments in category_to_segments.items():
|
||
if not acquire_lock(chapter_category):
|
||
logger.warning(
|
||
"章节锁竞争: category={}, 延迟重试",
|
||
chapter_category,
|
||
)
|
||
raise_retry()
|
||
|
||
try:
|
||
chapter, has_images = process_category(
|
||
chapter_category,
|
||
category_segments,
|
||
state,
|
||
user_profile,
|
||
user_birth_year,
|
||
llm,
|
||
)
|
||
if chapter and has_images:
|
||
chapters_to_enqueue.add(chapter.id)
|
||
finally:
|
||
release_lock(chapter_category)
|
||
|
||
return chapters_to_enqueue, len(segments)
|