search_transcript.py

"""Load a YouTube transcript and retrieve passages by keyword, without an LLM."""

import argparse
import json
import re
import sys

from langchain_core.documents import Document
from langchain_core.runnables import RunnableLambda

from capslane import CapslaneError
from capslane.langchain import CapslaneLoader


def search_passages(documents: list[Document], query: str) -> list[Document]:
    terms = set(re.findall(r"\w+", query.casefold()))
    if not terms:
        return []
    ranked = [(len(terms.intersection(re.findall(r"\w+", doc.page_content.casefold()))), index, doc) for index, doc in enumerate(documents)]
    ranked.sort(key=lambda item: (-item[0], item[1]))
    return [doc for score, _index, doc in ranked if score > 0][:3]


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("video", help="Public YouTube URL or video ID")
    parser.add_argument("--query", required=True, help="Words to find in the transcript")
    parser.add_argument("--lang", default=None)
    parser.add_argument("--mode", choices=["native", "auto", "generate"], default="native")
    parser.add_argument("--job", help="Resume a job for this same video")
    args = parser.parse_args(argv)
    loader = None
    try:
        loader = CapslaneLoader(args.video, lang=args.lang, mode=args.mode, job_id=args.job)
        documents = loader.load()
        retriever = RunnableLambda(lambda query: search_passages(documents, query))
        matches = retriever.invoke(args.query)
        print(json.dumps({"document_count": len(documents), "matches": [doc.model_dump() for doc in matches]}, ensure_ascii=False, indent=2))
        return 0
    except (CapslaneError, ValueError) as error:
        print(json.dumps({"error": getattr(error, "code", "invalid_configuration"), "message": str(error), "job_id": loader.job_id if loader else None, "request_id": getattr(error, "request_id", None)}), file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())