momentum-system/youtube_claims/recap/extract.py

"""Bundle -> structured recap record with the owner's local Ollama.

Pipeline per video: split the bundle into overlapping transcript chunks (only if it is longer
than the model comfortably handles) -> one structured-output call per chunk -> verify every
quote against the FULL transcript (snap to verbatim, drop what can't be found) -> merge the
chunk records (dedupe the overlap) -> normalise tickers. The verification + merge are pure and
unit-tested; only `ollama_chat` touches the network."""

from __future__ import annotations

import json
import re
import time
import urllib.error
import urllib.request
from collections import Counter

from youtube_claims.recap import prompt as P
from youtube_claims.recap import settings
from youtube_claims.recap.quotes import Transcript, _tok, transcript_text_from_bundle

_HEADER_MARK = "TRANSCRIPT (mm:ss text):"
_TICKER = re.compile(r"[A-Z]{1,5}")
_NULLISH = {"", "NULL", "NONE", "N/A", "NA", "UNKNOWN", "TBD", "?", "-"}


# ---------------------------------------------------------------- Ollama
def ollama_chat(system: str, user: str, *, model: str, host: str, schema: dict | None = None,
                num_ctx: int | None = None, think=None, timeout: int = 1800,
                temperature: float = 0.0, num_predict: int = 12000) -> tuple[str, dict]:
    """One non-streaming /api/chat call in structured-output mode. Returns (content, stats)."""
    payload = {"model": model, "stream": False, "format": schema or "json",
               "messages": [{"role": "system", "content": system}, {"role": "user", "content": user}],
               "options": {"temperature": temperature, "num_predict": num_predict}}
    if num_ctx:
        payload["options"]["num_ctx"] = num_ctx
    if think is not None:
        payload["think"] = think
    req = urllib.request.Request(host.rstrip("/") + "/api/chat", data=json.dumps(payload).encode(),
                                 headers={"Content-Type": "application/json"})
    t0 = time.time()
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            body = json.load(r)
    except urllib.error.HTTPError as e:     # surface Ollama's own message (bad model, bad think...)
        raise RuntimeError(f"ollama {e.code}: {e.read().decode(errors='replace')[:300]}") from e
    stats = {"seconds": round(time.time() - t0, 1), "prompt_tokens": body.get("prompt_eval_count"),
             "output_tokens": body.get("eval_count"), "done_reason": body.get("done_reason")}
    return (body.get("message") or {}).get("content") or "", stats


def parse_json(raw: str) -> dict:
    """Structured mode returns bare JSON; tolerate a stray code fence or preamble anyway."""
    raw = (raw or "").strip()
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        m = re.search(r"\{.*\}", raw, re.S)
        if not m:
            raise
        return json.loads(m.group(0))


# ---------------------------------------------------------------- chunking
def split_bundle(bundle: str) -> tuple[str, list[str]]:
    """(header incl. the transcript marker, transcript lines)."""
    if _HEADER_MARK not in bundle:
        return "", [ln for ln in bundle.splitlines() if ln.strip()]
    head, body = bundle.split(_HEADER_MARK, 1)
    return head + _HEADER_MARK, [ln for ln in body.splitlines() if ln.strip()]


def chunk_bundle(bundle: str, max_chars: int, overlap_lines: int = 4) -> list[str]:
    """The bundle itself if it fits, else header + consecutive transcript slices of about
    max_chars each, consecutive slices sharing `overlap_lines` lines so nothing said across
    a boundary is lost (the merge dedupes what both halves extracted)."""
    if len(bundle) <= max_chars:
        return [bundle]
    head, lines = split_bundle(bundle)
    budget = max(2000, max_chars - len(head))
    out, i = [], 0
    while i < len(lines):
        n, j = 0, i
        while j < len(lines) and (n + len(lines[j]) + 1 <= budget or j == i):
            n += len(lines[j]) + 1
            j += 1
        out.append(head + "\n" + "\n".join(lines[i:j]))
        if j >= len(lines):
            break
        i = max(i + 1, j - overlap_lines)
    return out


# ---------------------------------------------------------------- normalisation
def clean_ticker(raw) -> str | None:
    t = re.sub(r"[\s.$]", "", str(raw or "")).upper()
    if t in _NULLISH or not _TICKER.fullmatch(t):
        return None
    return t


def _num(x):
    if x is None or isinstance(x, bool):
        return None
    if isinstance(x, (int, float)):
        return float(x)
    m = re.search(r"-?\d[\d,]*\.?\d*", str(x))
    if not m:
        return None
    try:
        return float(m.group(0).replace(",", ""))
    except ValueError:
        return None


def _price_forms(p: float) -> set[str]:
    """How a spoken price shows up in a whisper transcript: 5.5 -> 5.50, 5.5, 550, $5.50;
    6.0 -> 6, 6.00, 600; 12.25 -> 12.25, 1225."""
    cents = int(round(p * 100))
    forms = {f"{p:.2f}", f"{p:.2f}".rstrip("0").rstrip("."), str(cents)}
    if cents % 100 == 0:
        forms.add(str(cents // 100))
    if p < 1:
        forms.add(str(cents))                     # "73 cents"
    return forms


def price_mentioned(tx: Transcript, p) -> bool:
    """Anti-hallucination for numbers: the price (in some spoken form) occurs in the transcript."""
    if p is None:
        return True
    toks = tx.token_set
    return any(f in toks for f in _price_forms(float(p)))


def name_mentioned(tx: Transcript, heard, guess) -> bool:
    """The stock was actually named: the heard name (as a phrase) or the guessed symbol occurs
    in the transcript, or the symbol spelled out letter by letter ("K I T T")."""
    ht = _tok(str(heard or ""))
    if ht and (" " + " ".join(ht) + " ") in tx.joined:
        return True
    g = (str(guess or "")).lower()
    return bool(g) and (g in tx.token_set or (" " + " ".join(g) + " ") in tx.joined)


def _enum(x, allowed, default=None):
    x = (str(x).strip().lower() if x is not None else None)
    return x if x in allowed else default


# ---------------------------------------------------------------- verification
MAX_QUOTE_WORDS = 50


def _snap(tx: Transcript, q) -> str | None:
    """The transcript's own words for quote q (None if not found). A model sometimes pastes a
    whole paragraph as the "quote"; it is verified in full, then cut to its first
    MAX_QUOTE_WORDS words -- still verbatim, and short enough to stay a pointer, not a copy."""
    s = tx.find(q) if isinstance(q, str) and q.strip() else None
    if s is not None:
        words = s.split()
        if len(words) > MAX_QUOTE_WORDS:
            s = " ".join(words[:MAX_QUOTE_WORDS])
    return s


def _quoted_list(tx: Transcript, items, key: str, stats: Counter, label: str) -> list[dict]:
    out = []
    for it in items or []:
        if not isinstance(it, dict):
            continue
        s = _snap(tx, it.get("quote"))
        if s is None:
            stats[f"dropped_{label}"] += 1
            continue
        out.append({key: str(it.get(key) or "").strip(), "quote": s})
    return out


def _descriptor_ok(tx: Transcript, s: str) -> bool:
    toks = _tok(s)
    if not toks:
        return False
    if len(toks) >= 3:
        return tx.find(s) is not None
    return (" " + " ".join(toks) + " ") in tx.joined


def verify_record(rec: dict, tx: Transcript, stats: Counter | None = None) -> dict:
    """Enforce the verbatim-quote rule in code. Quoted items whose quote is not in the
    transcript are dropped; accepted quotes are replaced by the transcript's exact words;
    a trade with no surviving quote at all (no evidence he described it) is dropped."""
    stats = stats if stats is not None else Counter()
    day_in = rec.get("day") or {}
    mq = _snap(tx, day_in.get("market_quote"))
    if day_in.get("market_quote") and mq is None:
        stats["dropped_market_quote"] += 1
    day = {
        "net_pnl": _num(day_in.get("net_pnl")),
        "green": day_in.get("green") if isinstance(day_in.get("green"), bool) else None,
        "market_temperature": _enum(day_in.get("market_temperature"), {"hot", "cold", "mixed"}),
        "market_quote": mq,
        "max_loss_hit": day_in.get("max_loss_hit") if isinstance(day_in.get("max_loss_hit"), bool) else None,
        "daily_goal_comment": day_in.get("daily_goal_comment") or None,
        "sizing_comment": day_in.get("sizing_comment") or None,
    }
    trades = []
    for t in rec.get("trades") or []:
        if not isinstance(t, dict):
            continue
        sc_in = t.get("stock_context") or {}
        cq = _snap(tx, sc_in.get("catalyst_quote"))
        sc = {"catalyst": (sc_in.get("catalyst") or None) if cq else None,
              "catalyst_quote": cq,
              "float": sc_in.get("float") or None,
              "pct_up_at_entry": sc_in.get("pct_up_at_entry") or None,
              "leading_gainer_rank": _num(sc_in.get("leading_gainer_rank")),
              "other": [s for s in (sc_in.get("other") or []) if isinstance(s, str) and _descriptor_ok(tx, s)]}
        out = {
            "ticker_heard": str(t.get("ticker_heard") or "").strip() or None,
            "ticker_guess": clean_ticker(t.get("ticker_guess")) or clean_ticker(t.get("ticker_heard")),
            "entry_time_hint": t.get("entry_time_hint") or None,
            "setup": _enum(t.get("setup"), set(P.SETUPS), "other"),
            "setup_quote": _snap(tx, t.get("setup_quote")),
            "entry_price": _num(t.get("entry_price")), "exit_price": _num(t.get("exit_price")),
            "stated_high": _num(t.get("stated_high")), "stated_low": _num(t.get("stated_low")),
            "shares": _num(t.get("shares")), "pnl": _num(t.get("pnl")),
            "outcome": _enum(t.get("outcome"), {"win", "loss", "scratch"}),
            "add_on": bool(t.get("add_on")),
            "stock_context": sc,
            "why_entered": _quoted_list(tx, t.get("why_entered"), "reason", stats, "why_entered"),
            "tape_reads": _quoted_list(tx, t.get("tape_reads"), "read", stats, "tape_reads"),
            "why_exited": _quoted_list(tx, t.get("why_exited"), "reason", stats, "why_exited"),
            "mistake_or_lesson": _quoted_list(tx, t.get("mistake_or_lesson"), "lesson", stats, "lessons"),
        }
        for k in ("entry_price", "exit_price", "stated_high", "stated_low"):
            if out[k] is not None and not price_mentioned(tx, out[k]):
                stats["dropped_unspoken_price"] += 1
                out[k] = None
        if t.get("setup_quote") and out["setup_quote"] is None:
            stats["dropped_setup_quote"] += 1
        if not (out["setup_quote"] or out["why_entered"] or out["why_exited"] or out["tape_reads"]
                or out["mistake_or_lesson"]):
            stats["dropped_trades"] += 1
            continue
        trades.append(out)
    passed = []
    for p in rec.get("passed_on") or []:
        if not isinstance(p, dict):
            continue
        s = _snap(tx, p.get("quote"))
        if s is None:
            stats["dropped_passed_on"] += 1
            continue
        if not name_mentioned(tx, p.get("ticker_heard"), p.get("ticker_guess")):
            stats["dropped_passed_unnamed"] += 1      # e.g. a scanner-list name he never said
            continue
        passed.append({"ticker_heard": str(p.get("ticker_heard") or "").strip() or None,
                       "ticker_guess": clean_ticker(p.get("ticker_guess")) or clean_ticker(p.get("ticker_heard")),
                       "reason": str(p.get("reason") or "").strip(), "quote": s})
    rules = _quoted_list(tx, rec.get("rules_stated"), "rule", stats, "rules")
    return {"video_kind": _enum(rec.get("video_kind"), set(P.KINDS), "other"),
            "session_date": rec.get("session_date") if _is_date(rec.get("session_date")) else None,
            "day": day, "trades": trades, "passed_on": passed, "rules_stated": rules}


def _is_date(s) -> bool:
    return isinstance(s, str) and bool(re.fullmatch(r"\d{4}-\d{2}-\d{2}", s))


# ---------------------------------------------------------------- merge
def _qkey(q) -> str:
    return " ".join(_tok(q or ""))


def _overlaps(a: str, b: str) -> bool:
    return bool(a) and bool(b) and (a in b or b in a)


def _same_trade(a: dict, b: dict) -> bool:
    if (a.get("ticker_guess") or a.get("ticker_heard")) != (b.get("ticker_guess") or b.get("ticker_heard")):
        return False
    qa = {_qkey(a.get("setup_quote"))} | {_qkey(x["quote"]) for k in ("why_entered", "why_exited") for x in a[k]}
    qb = {_qkey(b.get("setup_quote"))} | {_qkey(x["quote"]) for k in ("why_entered", "why_exited") for x in b[k]}
    qa.discard(""); qb.discard("")
    return any(_overlaps(x, y) for x in qa for y in qb)


def _union_quoted(a: list[dict], b: list[dict]) -> list[dict]:
    out = list(a)
    for x in b:
        if not any(_overlaps(_qkey(x["quote"]), _qkey(y["quote"])) for y in out):
            out.append(x)
    return out


def merge_records(parts: list[dict]) -> dict:
    """Merge verified per-chunk records into one. Chunks overlap, so the same trade / pass /
    rule can come back twice: trades are the same when ticker matches and a quote overlaps
    (their fields are then unioned, first non-null wins); passes dedupe by ticker; rules by
    quote. (A traded ticker may also be passed on -- he often skips a second entry.)"""
    parts = [p for p in parts if p]
    if not parts:
        return {"video_kind": "other", "session_date": None, "day": {}, "trades": [], "passed_on": [],
                "rules_stated": []}
    kinds = Counter(p.get("video_kind") for p in parts)
    kind = max(kinds, key=lambda k: (kinds[k], -[p.get("video_kind") for p in parts].index(k)))
    day: dict = {}
    for p in parts:
        for k, v in (p.get("day") or {}).items():
            if day.get(k) is None and v is not None:
                day[k] = v
    temps = Counter(p["day"].get("market_temperature") for p in parts if (p.get("day") or {}).get("market_temperature"))
    if len(temps) > 1:
        top = temps.most_common()
        day["market_temperature"] = top[0][0] if top[0][1] > top[1][1] else "mixed"
    trades: list[dict] = []
    for p in parts:
        for t in p.get("trades") or []:
            dup = next((x for x in trades if _same_trade(x, t)), None)
            if dup is None:
                trades.append(json.loads(json.dumps(t)))
                continue
            for k, v in t.items():
                if k in ("why_entered", "tape_reads", "why_exited", "mistake_or_lesson"):
                    dup[k] = _union_quoted(dup[k], v)
                elif k == "stock_context":
                    for ck, cv in (v or {}).items():
                        if ck == "other":
                            dup[k]["other"] = list(dict.fromkeys((dup[k].get("other") or []) + (cv or [])))
                        elif dup[k].get(ck) is None:
                            dup[k][ck] = cv
                elif dup.get(k) is None and v is not None:
                    dup[k] = v
    passed, seen = [], set()
    for p in parts:
        for x in p.get("passed_on") or []:
            key = x.get("ticker_guess") or (x.get("ticker_heard") or "").upper()
            if key in seen:
                continue
            seen.add(key)
            passed.append(x)
    rules: list[dict] = []
    for p in parts:
        rules = _union_quoted(rules, p.get("rules_stated") or [])
    return {"video_kind": kind, "session_date": next((p["session_date"] for p in parts if p.get("session_date")), None),
            "day": day, "trades": trades, "passed_on": passed, "rules_stated": rules}


# ---------------------------------------------------------------- spoken prices
def scanner_prices(bundle: str) -> dict[str, float]:
    """{symbol: price} from the bundle's OUR SCANNER LIST lines ('SYM 12.3% $4.5')."""
    out: dict[str, float] = {}
    head = bundle.split(_HEADER_MARK, 1)[0]
    for sym, px in re.findall(r"\b([A-Z]{1,5}) -?[\d.]+% \$([\d.]+)", head):
        out[sym] = max(out.get(sym, 0.0), float(px))
    return out


def fix_spoken_prices(rec: dict, prices: dict[str, float]) -> int:
    """ASR/LLM render "six ninety-nine" as 699 -- Ross trades small caps, so a trade price
    >= 100 is divided by 100 unless the scanner saw that ticker above $50. Returns #fixes."""
    n = 0
    for t in rec.get("trades") or []:
        if prices.get(t.get("ticker_guess") or "", 0) >= 50:
            continue
        for k in ("entry_price", "exit_price", "stated_high", "stated_low"):
            v = t.get(k)
            if v is not None and v >= 100:
                t[k] = round(v / 100, 4)
                n += 1
    return n


# ---------------------------------------------------------------- driver
def bundle_meta(bundle: str) -> dict:
    meta = {}
    for ln in bundle.splitlines()[:6]:
        k, _, v = ln.partition(": ")
        if k in ("video_id", "upload_date", "title", "transcript_source"):
            meta[k] = v.strip()
    return meta


def extract_bundle(bundle: str, *, model: str | None = None, host: str | None = None,
                   num_ctx: int | None = None, think=None, chunk_chars: int | None = None,
                   system: str | None = None, split: bool | None = None, log=print) -> tuple[dict, dict]:
    """Extract one video. Returns (record per EXTRACT_SPEC, run metadata incl. drop stats).
    split=True asks for trades/day and for passed-on/rules in two focused calls per chunk
    (better recall from a small local model at ~2x the time)."""
    split = settings.SPLIT_CALLS if split is None else split
    think = settings.OLLAMA_THINK if think is None else think
    from youtube_claims.recap.llm import Endpoint, OllamaPool, from_settings
    if host:                                   # an explicit endpoint (evaluation runs)
        pool = OllamaPool([Endpoint(host.rstrip("/"))], model or settings.OLLAMA_MODEL,
                          settings.OLLAMA_NUM_CTX if num_ctx is None else (num_ctx or None), think,
                          call=ollama_chat)
    else:
        pool = from_settings()
        if model:
            pool.default_model = model
        if num_ctx is not None:
            pool.default_ctx = num_ctx or None
        pool.default_think = think
    model = pool.default_model
    chunk_chars = chunk_chars or settings.CHUNK_CHARS
    meta = bundle_meta(bundle)
    tx = Transcript(transcript_text_from_bundle(bundle))
    chunks = chunk_bundle(bundle, chunk_chars)
    stats: Counter = Counter()
    parts, calls = [], []
    tasks = P.SPLIT_TASKS if split else [(None, P.SCHEMA)]
    for i, ch in enumerate(chunks, 1):
        note = P.PART_NOTE.format(i=i, n=len(chunks)) if len(chunks) > 1 else ""
        data: dict | None = {}
        for task, schema in tasks:
            user = P.USER_TMPL.format(part_note=note + (task or ""), bundle=ch)
            got = None
            for attempt in range(2):
                try:
                    raw, st = pool.chat(system or P.SYSTEM, user, schema=schema, timeout=settings.OLLAMA_TIMEOUT,
                                        temperature=0.0 if attempt == 0 else 0.3,
                                        num_predict=settings.OLLAMA_NUM_PREDICT)
                    calls.append(st)
                    got = parse_json(raw)
                    break
                except (json.JSONDecodeError, ValueError) as e:
                    stats["json_errors"] += 1
                    log(f"  chunk {i}/{len(chunks)} bad JSON ({e}); retrying" if attempt == 0 else
                        f"  chunk {i}/{len(chunks)} bad JSON again; skipping")
            if got is None:
                data = None
                break
            data.update({k: v for k, v in got.items() if k in schema["properties"]})  # a call may echo other keys empty
        if data is None:
            stats["chunks_failed"] += 1
            continue
        for k in ("trades", "passed_on", "rules_stated"):
            stats[f"raw_{k}"] += len(data.get(k) or [])
        parts.append(verify_record(data, tx, stats))
    if not parts:
        raise RuntimeError(f"extraction produced nothing for {meta.get('video_id')} ({dict(stats)})")
    rec = merge_records(parts)
    fixed = fix_spoken_prices(rec, scanner_prices(bundle))
    if fixed:
        stats["prices_rescaled"] += fixed
    up = meta.get("upload_date") or ""
    out = {"video_id": meta.get("video_id"), "upload_date": up,
           "session_date": rec.get("session_date"), "video_kind": rec["video_kind"],
           "day": rec["day"], "trades": rec["trades"], "passed_on": rec["passed_on"],
           "rules_stated": rec["rules_stated"]}
    run = {"model": model, "hosts": sorted({c.get("host") for c in calls if c.get("host")}),
           "chunks": len(chunks), "split": split, "calls": calls,
           "seconds": round(sum(c["seconds"] for c in calls), 1), "stats": dict(stats)}
    return out, run