src/pipeline.py

"""Generation pipeline shared by UI, scripts, and tests."""

from __future__ import annotations

from datetime import datetime
from html import escape
from pathlib import Path

from src.config import TRACE_DIR, get_runtime_settings, runtime_status
from src.models.llama_cpp_runner import (
    generate_persona_and_diary,
    get_text_runtime_fallbacks,
    reset_text_runtime_fallbacks,
)
from src.models.schema import GenerationResult
from src.models.vision_runner import VisionRunResult, understand_object_with_metadata
from src.traces.logger import build_trace, save_trace


def generate_object_diary(
    image_path: str | None,
    description: str,
    mode: str,
    *,
    trace_dir: Path = TRACE_DIR,
    save: bool = True,
    trace_id: str | None = None,
    created_at: datetime | None = None,
) -> GenerationResult:
    settings = get_runtime_settings()
    vision_result = understand_object_with_metadata(image_path, description, settings=settings)
    object_understanding = vision_result.object_understanding
    reset_text_runtime_fallbacks()
    persona, diary = generate_persona_and_diary(object_understanding, mode)
    text_fallbacks = get_text_runtime_fallbacks()
    trace = build_trace(
        image_path=image_path,
        description=description,
        mode=mode,
        object_understanding=object_understanding,
        persona=persona,
        diary=diary,
        trace_id=trace_id,
        created_at=created_at,
        model_runtime=runtime_status(settings),
        fallbacks=_runtime_fallbacks(
            settings.vision_backend,
            settings.text_backend,
            vision_result,
            text_fallbacks,
        ),
    )
    trace_path = save_trace(trace, trace_dir) if save else ""

    return GenerationResult(
        object_understanding=object_understanding,
        persona=persona,
        diary=diary,
        trace=trace,
        trace_path=trace_path,
    )


def format_diary_markdown(title: str, english: str, chinese: str) -> str:
    return f"""
<section class="diary-entry">
  <h2>{escape(title)}</h2>
  <p>{escape(english)}</p>
  <section class="lang-zh block zh-helper">
    <strong>中文辅助</strong>
    <p>{escape(chinese)}</p>
  </section>
</section>
"""


def _runtime_fallbacks(
    vision_backend: str,
    text_backend: str,
    vision_result: VisionRunResult,
    text_fallbacks: list[str] | None = None,
) -> list[str]:
    clean_vision_backend = vision_backend.strip().lower()
    clean_text_backend = text_backend.strip().lower()
    if clean_vision_backend == "mock" and clean_text_backend == "mock":
        return ["mock-runtime"]

    fallbacks = list(vision_result.fallbacks)
    for marker in text_fallbacks or []:
        if marker not in fallbacks:
            fallbacks.append(marker)
    if clean_vision_backend == "mock":
        fallbacks.append("mock-vision-runtime")
    if clean_text_backend == "mock":
        fallbacks.append("mock-text-runtime")
    return fallbacks