second-brain/src/second_brain/main.py

"""
second-brain CLI entry point.

Commands:
  add       — queue a source URL (video or article)
  process   — run the full pipeline on pending/pulled/transcribed sources
  serve     — launch the FastAPI web UI
  compile   — run the wiki compiler on accepted sources
  schedule  — start the overnight scheduler
"""

from __future__ import annotations

import sys
from typing import Optional

import click

from second_brain.config import DOMAINS, load_config
from second_brain.database import get_database
from second_brain.models import Source, SourceStatus, SourceType, utcnow


# ---------------------------------------------------------------------------
# CLI group
# ---------------------------------------------------------------------------


@click.group()
def cli() -> None:
    """second-brain — personal knowledge extraction pipeline."""


# ---------------------------------------------------------------------------
# add
# ---------------------------------------------------------------------------


@cli.command()
@click.argument("url")
@click.option(
    "--domain",
    "-d",
    type=click.Choice(DOMAINS),
    default="development",
    show_default=True,
    help="Knowledge domain for extraction.",
)
@click.option("--focus", "-f", default=None, help="Optional focus directive.")
@click.option("--title", "-t", default=None, help="Override title.")
def add(url: str, domain: str, focus: Optional[str], title: Optional[str]) -> None:
    """Queue a source URL for processing.

    Pasting a YouTube playlist URL (the `/playlist?list=…` form) auto-
    expands into one source row per video — no flag, just paste.
    `watch?v=…&list=…` is treated as a single video (use the canonical
    playlist URL to fan out).
    """
    from second_brain.sources_service import (
        DEFAULT_PLAYLIST_MAX_ITEMS,
        add_playlist,
        add_source,
        is_youtube_playlist_url,
    )

    config = load_config()
    config.ensure_dirs()
    db = get_database()

    # --- playlist fan-out path ---------------------------------------------
    if is_youtube_playlist_url(url):
        try:
            with db.session() as sess:
                result = add_playlist(
                    sess,
                    url=url,
                    domain=domain,
                    focus=focus,
                    max_items=DEFAULT_PLAYLIST_MAX_ITEMS,
                )
        except ValueError as exc:
            click.echo(f"[add] Error: {exc}", err=True)
            sys.exit(2)

        click.echo(
            f"[add] Playlist expanded: {result.expanded} entr"
            f"{'y' if result.expanded == 1 else 'ies'} "
            f"(cap {DEFAULT_PLAYLIST_MAX_ITEMS})"
        )
        click.echo(
            f"      added={result.added}  duplicates={result.duplicates}  "
            f"failed={result.failed}"
        )
        for u, why in result.failures[:5]:
            click.echo(f"      ! {u}: {why}", err=True)
        return

    # --- single-source path ------------------------------------------------
    try:
        with db.session() as sess:
            result = add_source(
                sess, url=url, domain=domain, focus=focus, title=title
            )
            # Capture display fields inside the session — `result.source`
            # detaches on commit and accessing it later would error.
            source_id = result.source.id
            source_url = result.source.url
            source_type_value = result.source.source_type.value
            existing_status = result.source.status.value
            added = result.added
    except ValueError as exc:
        click.echo(f"[add] Error: {exc}", err=True)
        sys.exit(2)

    if not added:
        click.echo(f"[add] Already queued (id={source_id}, status={existing_status})")
        return

    click.echo(f"[add] Queued {source_type_value}: {source_url}")
    click.echo(f"      domain={domain}" + (f"  focus={focus}" if focus else ""))


# ---------------------------------------------------------------------------
# process
# ---------------------------------------------------------------------------


@cli.command()
@click.option("--limit", "-n", default=None, type=int, help="Max sources to process.")
@click.option(
    "--skip-pull", is_flag=True, default=False, help="Skip the pull/download step."
)
@click.option(
    "--skip-transcribe",
    is_flag=True,
    default=False,
    help="Skip the transcription step.",
)
@click.option(
    "--skip-extract", is_flag=True, default=False, help="Skip the LLM extraction step."
)
def process(
    limit: Optional[int],
    skip_pull: bool,
    skip_transcribe: bool,
    skip_extract: bool,
) -> None:
    """Run the dev-side pipeline: article pulls + extract/embed.

    Video pull + transcribe have moved to the tower-side
    `transcribe-worker` daemon (faster-whisper on the 3080); this command
    is article-only on the pull side, but still picks up TRANSCRIBED rows
    of any source_type for the extract/embed step.
    """
    from second_brain.adapters.article import ArticleAdapter
    from second_brain.context.assembler import ContextAssembler
    from second_brain.extractor.engine import ExtractionEngine
    from second_brain.models import Extraction

    config = load_config()
    config.ensure_dirs()
    db = get_database()

    art_adapter = ArticleAdapter(config)
    assembler = ContextAssembler(config)

    engine: Optional[ExtractionEngine] = None
    if not skip_extract:
        ex_cfg = config.extractor
        try:
            engine = ExtractionEngine(
                backend=ex_cfg["backend"],
                model=ex_cfg["model"],
                api_key=config.anthropic_api_key,
                cli_binary=ex_cfg["cli_binary"],
                timeout_sec=ex_cfg["timeout_sec"],
            )
        except ValueError as exc:
            click.echo(f"[process] Warning: {exc}")
            click.echo("[process] Extraction step will be skipped.")
            skip_extract = True

    # Dev side: pull only ARTICLE sources (videos belong to the tower
    # worker) — but pick up TRANSCRIBED of any type for extract/embed.
    with db.session() as sess:
        pull_ids = []
        if not skip_pull:
            pull_ids = [
                row[0]
                for row in sess.query(Source.id)
                .filter(
                    Source.status == SourceStatus.PENDING,
                    Source.source_type == SourceType.ARTICLE,
                )
                .all()
            ]
        extract_ids = [
            row[0]
            for row in sess.query(Source.id)
            .filter(Source.status == SourceStatus.TRANSCRIBED)
            .all()
        ]
        # Keep order stable (pulls first so freshly-transcribed articles
        # roll into extract in the same invocation) and de-duplicate.
        source_ids: list[int] = []
        for sid in pull_ids + extract_ids:
            if sid not in source_ids:
                source_ids.append(sid)
        if limit:
            source_ids = source_ids[:limit]

    if not source_ids:
        click.echo("[process] No sources to process.")
        return

    click.echo(f"[process] Processing {len(source_ids)} source(s)…")
    if skip_transcribe:
        click.echo("[process] (--skip-transcribe is now a no-op — transcription "
                   "runs on the tower worker.)")

    for source_id in source_ids:
        with db.session() as sess:
            source = sess.get(Source, source_id)
            click.echo(f"\n  [{source.id}] {source.url[:80]}")

            # Pull step — articles only on the dev side.
            if (
                not skip_pull
                and source.status == SourceStatus.PENDING
                and source.source_type == SourceType.ARTICLE
            ):
                click.echo("    → pull (article)")
                ok = art_adapter.pull(source)
                if not ok:
                    click.echo(f"    ✗ pull failed: {source.error_message}")
                    continue

            # Extract step
            if not skip_extract and source.status == SourceStatus.TRANSCRIBED and engine:
                click.echo("    → extract")
                try:
                    prompt = assembler.assemble(source)
                    result = engine.extract(prompt, source)

                    existing_ext = sess.query(Extraction).filter_by(
                        source_id=source.id
                    ).first()
                    if existing_ext:
                        for k, v in result.to_db_dict().items():
                            setattr(existing_ext, k, v)
                        existing_ext.raw_response = getattr(result, "_raw_response", None)
                        existing_ext.input_tokens = getattr(result, "_input_tokens", None)
                        existing_ext.output_tokens = getattr(result, "_output_tokens", None)
                        ext_row = existing_ext
                    else:
                        ext_row = Extraction(
                            source_id=source.id,
                            raw_response=getattr(result, "_raw_response", None),
                            input_tokens=getattr(result, "_input_tokens", None),
                            output_tokens=getattr(result, "_output_tokens", None),
                            **result.to_db_dict(),
                        )
                        sess.add(ext_row)

                    source.status = SourceStatus.ANALYZED
                    source.updated_at = utcnow()
                    # Flush so the new extraction gets its PK before we kick
                    # off embedding (we key public.embeddings rows by it).
                    sess.flush()
                    extraction_id = ext_row.id
                    extraction_summary = ext_row.summary
                    click.echo("    ✓ analyzed")

                    if config.embeddings.get("enabled", True):
                        from second_brain.embeddings import embed_extraction

                        n = embed_extraction(extraction_id, extraction_summary)
                        if n is not None:
                            click.echo(f"    ✓ embedded ({n} chunk(s))")
                except Exception as exc:
                    source.error_message = str(exc)
                    source.status = SourceStatus.FAILED
                    click.echo(f"    ✗ extraction failed: {exc}")

    click.echo("\n[process] Done.")


# ---------------------------------------------------------------------------
# serve
# ---------------------------------------------------------------------------


@cli.command()
@click.option("--host", default="127.0.0.1", show_default=True)
@click.option("--port", default=8000, show_default=True, type=int)
@click.option("--reload", is_flag=True, default=False)
def serve(host: str, port: int, reload: bool) -> None:
    """Launch the web UI."""
    import uvicorn

    click.echo(f"[serve] Starting on http://{host}:{port}")
    uvicorn.run(
        "second_brain.web.app:app",
        host=host,
        port=port,
        reload=reload,
    )


# ---------------------------------------------------------------------------
# compile
# ---------------------------------------------------------------------------


@cli.command("compile")
def compile_wiki() -> None:
    """Run the wiki compiler on accepted sources."""
    from second_brain.compiler.wiki import WikiCompiler

    config = load_config()
    compiler = WikiCompiler(config)
    stats = compiler.run()
    click.echo(
        f"[compile] processed={stats['processed']} "
        f"created={stats['created']} updated={stats['updated']} "
        f"errors={stats['errors']}"
    )


# ---------------------------------------------------------------------------
# schedule
# ---------------------------------------------------------------------------


@cli.command()
@click.option("--dry-run", is_flag=True, default=False)
def schedule(dry_run: bool) -> None:
    """Start the overnight scheduler."""
    from second_brain.scheduler.runner import Scheduler

    config = load_config()
    scheduler = Scheduler(config)
    scheduler.run(dry_run=dry_run)


# ---------------------------------------------------------------------------
# transcribe-worker (tower-side daemon)
# ---------------------------------------------------------------------------


@cli.command("transcribe-worker")
@click.option("--once", is_flag=True, default=False,
              help="Process at most one source then exit (useful for ad-hoc runs / tests).")
@click.option("--poll-interval", default=None, type=int,
              help="Seconds between empty-queue polls (default 15).")
@click.option("--worker-id", default=None,
              help="Identifier stamped into sources.claimed_by (default tower:<hostname>).")
def transcribe_worker(
    once: bool,
    poll_interval: Optional[int],
    worker_id: Optional[str],
) -> None:
    """Long-running tower-side worker: claim → yt-dlp pull → faster-whisper.

    Reads pipeline_settings each iteration so the dashboard's enable
    toggle / active-hours window / max-items-per-run / max-video-length
    all take effect within one poll cycle.
    """
    import logging

    from second_brain.scheduler.transcribe_worker import (
        DEFAULT_POLL_INTERVAL_SECONDS,
        TranscribeWorker,
    )

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )

    config = load_config()
    config.ensure_dirs()
    worker = TranscribeWorker(
        config,
        worker_id=worker_id,
        poll_interval=poll_interval or DEFAULT_POLL_INTERVAL_SECONDS,
    )
    processed = worker.run(once=once)
    click.echo(f"[transcribe-worker] exiting; processed={processed}")


if __name__ == "__main__":
    cli()