main.py

"""A long-lived process that maintains news embeddings.

It scans `ingest_stage_state` for pending work and runs the `index` stage
itself. Nothing pushes to it: whatever collects articles writes `news_items`
plus a `fetch` stage row, and this worker picks them up on its next pass, so
the two sides share only Postgres.

Each stage is idempotent and records its own `ingest_stage_state`, so a tick is
simply "drain whatever is pending" and a crash mid-tick leaves the remainder
for the next pass. Ticks never overlap — the wait happens after a scan ends.
"""

from dotenv import load_dotenv

load_dotenv()  # must run before modules that read env at import (db.py)

import time
from typing import Callable
from health import WorkerState, start_health_server

from src.config import Config, load_config
from src.work import pending_index
from src.indexer import index_news_item
from src.embed import embed_texts
from src.embed_server import start_embed_server

# The worker runs a single stage: index (whole-doc + sentence-chunk embeddings).
# Nothing is derived beyond the vectors themselves — no entity extraction, no
# classification, no LLM in the loop. Matching happens at query time against
# these embeddings.


def process_stage(name: str, ids: list[str], run: Callable[[str], None],
                  state: WorkerState) -> None:
    for news_id in ids:
        # Heartbeat before each item so a long backfill tick keeps reporting
        # liveness (health.py reads last_progress_at) instead of going stale.
        state.last_progress_at = time.time()
        try:
            run(news_id)
            state.processed[name] += 1
        except Exception as err:
            # The stage has already recorded its own 'failed' row (and will be
            # retried next pass); log and keep draining so one bad item can't
            # stall the others.
            print(f"[algorithm] {name} failed newsItemId={news_id}: {err}")


def tick(config: Config, state: WorkerState) -> None:
    process_stage("index", pending_index(config.batch_size),
                  index_news_item, state)


def main() -> None:
    config = load_config()
    print(
        f"[algorithm] starting; tick={config.tick_seconds}s "
        f"batch={config.batch_size} (index-only)"
    )

    state = WorkerState()
    start_health_server(config.health_port, state, config.tick_seconds)
    # Text->vector endpoint, sharing this worker's already-resident model.
    # No-ops unless EMBED_SERVICE_SECRET is set.
    start_embed_server(config.embed_port, config.embed_secret)

    # Warm bge-m3 up front so the first indexing tick doesn't block on the ~2GB
    # model load. The load itself can outlast the boot-grace window (cold import
    # of torch/transformers + reading the weights from disk), so seed the
    # progress heartbeat *before* it — health then reads the worker as alive
    # throughout the load instead of tripping when grace expires. Best-effort.
    state.last_progress_at = time.time()
    try:
        embed_texts(["warmup"])
        state.last_progress_at = time.time()
        print("[algorithm] bge-m3 warm")
    except Exception as err:
        print(f"[algorithm] warm skipped: {err}")

    while True:
        started_at = time.time()
        try:
            tick(config, state)
            state.last_tick_at = time.time()
            state.last_tick_error = None
        except Exception as err:
            state.last_tick_error = str(err)
            print(f"[algorithm] tick failed: {err}")
        state.ticks += 1
        time.sleep(max(0.0, config.tick_seconds - (time.time() - started_at)))


if __name__ == "__main__":
    main()