diff --git a/tests/unit/test_youtube_extraction.py b/tests/unit/test_youtube_extraction.py index 3d2d90c..fcdc412 100644 --- a/tests/unit/test_youtube_extraction.py +++ b/tests/unit/test_youtube_extraction.py @@ -1,7 +1,8 @@ """Pure-helper tests for the claim extractor — the anti-hallucination quote match + chunking. The Ollama call itself is integration-tested live (not mocked here).""" -from youtube_claims.extraction import _chunks, _norm, _normalize, _quote_supported +from youtube_claims.extraction import ( + _chunks, _clean_ticker, _norm, _normalize, _quote_supported) def test_norm_collapses_whitespace(): @@ -26,6 +27,23 @@ def test_normalize_drops_claim_without_verbatim_quote(): assert ok and ok["asset_ticker"] == "AAPL" and ok["timestamp_start"] == 0 +def test_clean_ticker_coerces_nullish_and_spacing(): + assert _clean_ticker("NULL") is None and _clean_ticker("none") is None + assert _clean_ticker("$AAPL") == "AAPL" + assert _clean_ticker("B VC") == "BVC" # ASR space-mangling + assert _clean_ticker("") is None + + +def test_asset_specific_filter_drops_generic_commentary(): + # a "NULL"-ticker generic tutorial line with no asset is dropped + generic = {"asset_ticker": "NULL", "asset_name": None, "verbatim_quote": "buy the dip"} + assert _normalize(generic, 0, 6) is None + # a real asset claim survives, ticker cleaned + real = {"asset_ticker": "null", "asset_name": "Apple", "verbatim_quote": "apple broke out"} + keep = _normalize(real, 0, 6) + assert keep and keep["asset_ticker"] is None and keep["asset_name"] == "Apple" + + def test_chunks_group_segments_with_time_span(): segs = [{"start": float(i), "end": float(i + 1), "text": "x" * 100} for i in range(200)] wins = list(_chunks(segs, max_chars=1000)) diff --git a/youtube_claims/extraction.py b/youtube_claims/extraction.py index d806bb2..6e9301c 100644 --- a/youtube_claims/extraction.py +++ b/youtube_claims/extraction.py @@ -80,14 +80,35 @@ def _chunks(segments: list[dict], max_chars: int = 8000): yield buf, start, buf[-1][0].get("end") -def _normalize(c: dict, win_start, win_end) -> dict | None: +_NULLISH = {"", "NULL", "NONE", "N/A", "NA", "-", "--", "UNKNOWN", "TBD", "IT", "THE STOCK"} + + +def _clean_ticker(raw) -> str | None: + """The LLM sometimes returns the literal string 'NULL' or a space-mangled ASR ticker + ('B VC'); coerce those to a real None / contiguous symbol.""" + t = (raw or "").strip().upper().lstrip("$") + if t in _NULLISH: + return None + t = t.replace(" ", "") # 'B VC' -> 'BVC' (tickers are single tokens) + return t or None + + +def _clean_name(raw) -> str | None: + n = (raw or "").strip() + return None if (not n or n.upper() in _NULLISH) else n + + +def _normalize(c: dict, win_start, win_end, require_asset: bool = True) -> dict | None: vq = (c.get("verbatim_quote") or "").strip() if not vq: return None # mandatory — drop unsupported claims - tkr = (c.get("asset_ticker") or "").strip().upper() or None + tkr = _clean_ticker(c.get("asset_ticker")) + name = _clean_name(c.get("asset_name")) + if require_asset and not tkr and not name: + return None # not an asset-specific claim (generic commentary) return { "asset_ticker": tkr, - "asset_name": (c.get("asset_name") or None), + "asset_name": name, "asset_class": (c.get("asset_class") or None), "direction": (c.get("direction") or None), "claim_text": (c.get("claim_text") or None),