pipeline/orchestrator.py

"""End-to-end orchestration of the blog-post pipeline.

Kept UI-agnostic: `run()` is a generator that yields (progress_fraction, message,
partial_result) tuples so any front-end can drive a progress bar. The final yield
carries the complete result dict.
"""
from __future__ import annotations

from pathlib import Path
from typing import Iterator, Tuple

from . import (
    cache,
    captions,
    config,
    extract,
    images,
    llm,
    openpagerank,
    search_terms,
    searxng_client,
    writer,
)


def run(
    hf_token: str,
    topic: str,
    primary_keyword: str,
    secondary_keyword: str,
    brief: str,
    target_wordcount: int = config.DEFAULT_WORD_COUNT,
    content_goal: str = writer.DEFAULT_GOAL,
) -> Iterator[Tuple[float, str, dict]]:
    if not topic.strip():
        raise ValueError("Please enter a topic.")
    client = llm.make_client(hf_token)  # raises if token missing → billed to user

    try:
        target_wordcount = int(target_wordcount)
    except (TypeError, ValueError):
        target_wordcount = config.DEFAULT_WORD_COUNT
    target_wordcount = max(300, min(target_wordcount, 5000))
    if content_goal not in writer.GOAL_GUIDANCE:
        content_goal = writer.DEFAULT_GOAL

    key = cache.run_key(
        topic, primary_keyword, secondary_keyword, brief, str(target_wordcount), content_goal
    )
    run_dir = config.OUT_DIR / key
    run_dir.mkdir(parents=True, exist_ok=True)
    result: dict = {"key": key}

    # 1) search terms
    yield 0.05, "Generating search terms…", result
    terms = cache.get(key, "terms") or search_terms.generate_search_terms(
        client, topic, primary_keyword, secondary_keyword, brief
    )
    cache.put(key, "terms", terms)
    result["terms"] = terms

    # 2) SearXNG search
    yield 0.15, f"Searching the web for {len(terms)} queries…", result
    candidates = cache.get(key, "candidates") or searxng_client.search(terms, config.TOP_N)
    if not candidates:
        raise RuntimeError(
            "No search results from SearXNG. Is the SearXNG service running "
            "(check the Space logs)?"
        )
    cache.put(key, "candidates", candidates)
    result["candidates"] = candidates

    # 3) OpenPageRank authority ranking
    yield 0.30, f"Ranking {len(candidates)} domains by OpenPageRank…", result
    top = openpagerank.rank_results(candidates, config.TOP_K)
    cache.put(key, "top", top)
    result["top"] = top

    # 4) extract source material
    yield 0.40, f"Extracting content from top {len(top)} pages…", result
    sources = extract.extract_sources(top)
    if not sources:
        raise RuntimeError("Could not extract readable content from the top pages.")
    result["sources"] = [{k: v for k, v in s.items() if k != "text"} for s in sources]

    # 5) write the post
    yield 0.55, f"Writing the ~{target_wordcount}-word {content_goal} blog post…", result
    markdown = writer.write_post(
        client, topic, primary_keyword, secondary_keyword, brief, sources,
        target_wordcount=target_wordcount, content_goal=content_goal,
    )
    (run_dir / "post.md").write_text(markdown, encoding="utf-8")
    result["markdown"] = markdown

    # 6) generate images (remote FLUX.1-schnell inference calls, billed to the user)
    scenes = writer.parse_image_markers(markdown)
    yield 0.70, f"Generating {len(scenes)} images with FLUX.1-schnell…", result
    imgs = images.generate_images(client, hf_token, topic, scenes, run_dir)
    result["images"] = imgs
    n_ok = sum(1 for im in imgs if im.get("path"))
    if n_ok == 0 and imgs:
        yield 0.72, f"⚠ Image generation failed: {imgs[0].get('error', 'unknown error')}", result
    else:
        yield 0.72, f"Generated {n_ok}/{len(imgs)} images.", result

    # 7) caption images
    yield 0.85, "Captioning images…", result
    imgs = captions.caption_images(client, imgs)
    result["images"] = imgs

    # 8) build docx
    yield 0.95, "Building the .docx file…", result
    from . import docx_builder

    docx_path = docx_builder.build_docx(markdown, imgs, run_dir / "blog.docx")
    result["docx_path"] = str(docx_path)

    yield 1.0, "Done.", result