fithia2/scripts/llm_parser_test.py

"""LLM vs Rule-Based Parser Comparison Test.

qwen3.5:9b (Ollama local) vs RuleBasedParser on real SEC exhibit texts.
Tests: direction accuracy, guidance detection, EPS extraction, low-confidence rescue.

Usage:
    uv run python3 scripts/llm_parser_test.py
"""
from __future__ import annotations

import json
import os
import sys
import time
from pathlib import Path

import pandas as pd
import requests

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from libs.common.file_store import read_exhibit
from libs.parser.rule_parser import RuleBasedParser
from libs.parser.text_normalizer import looks_like_html, normalize_text

OLLAMA_URL = "http://localhost:11434/api/generate"
MODEL = "gpt-oss:20b"
EXHIBIT_CACHE = "data/cache/exhibits"

_rule_parser = RuleBasedParser()


# ── LLM Prompt ────────────────────────────────────────────────────────────────

SYSTEM_PROMPT = """You are a financial document analyst specializing in SEC 8-K press releases.
Analyze the given earnings/corporate announcement text and extract structured information.
Respond ONLY with a valid JSON object, no markdown, no explanation."""

USER_PROMPT_TEMPLATE = """Analyze this SEC press release and return ONLY a JSON object with these fields:

{{
  "event_type": "earnings_release|guidance_update|material_contract|management_change|other_material_event|unknown",
  "event_direction": "bullish|bearish|neutral|unknown",
  "guidance_status": "raised|lowered|maintained|withdrawn|not_provided",
  "eps_actual": <float or null>,
  "eps_estimate": <float or null>,
  "eps_beat_pct": <float or null, positive=beat, negative=miss>,
  "revenue_actual_bn": <float or null, in billions>,
  "revenue_beat_pct": <float or null>,
  "key_signal": "<one sentence: the single most important signal>",
  "confidence": <0.0-1.0>
}}

Rules:
- event_direction "bullish" = results/news clearly positive for stock price
- event_direction "bearish" = results/news clearly negative
- guidance_status "raised" = company raised future guidance/outlook
- Extract EPS numbers only if explicitly stated (e.g. "$2.34 vs $2.10 estimate")
- confidence reflects how clearly the document supports your answer

DOCUMENT:
{text}"""


# ── Helpers ───────────────────────────────────────────────────────────────────

def extract_accession_no(event_id: str) -> str | None:
    parts = event_id.split("::")
    if len(parts) >= 7:
        return parts[6]
    return None


def get_exhibit_text(accession_no: str) -> str | None:
    try:
        text = read_exhibit(accession_no, "EX-99.1")
        normalized = normalize_text(text, is_html=looks_like_html(text))
        # Truncate to ~3000 chars for LLM (key info is in first part)
        return normalized[:3000]
    except FileNotFoundError:
        return None


def run_rule_parser(accession_no: str, form_type: str = "8-K") -> dict:
    try:
        text = read_exhibit(accession_no, "EX-99.1")
        normalized = normalize_text(text, is_html=looks_like_html(text))
        output = _rule_parser.parse(
            document_id=accession_no,
            form_type=form_type,
            text=normalized,
            metadata={"item_numbers": []},
        )
        return {
            "event_type": output.event_type,
            "event_direction": output.event_direction,
            "guidance_status": output.guidance.status if output.guidance else "not_provided",
            "confidence": output.confidence.overall,
            "signal_strength": output.signals.demand_strength if output.signals else None,
        }
    except Exception as e:
        return {"error": str(e)}


def run_llm_parser(text: str) -> dict:
    prompt = USER_PROMPT_TEMPLATE.format(text=text)
    payload = {
        "model": MODEL,
        "prompt": prompt,
        "system": SYSTEM_PROMPT,
        "stream": False,
        "options": {"temperature": 0.1, "num_predict": 400},
    }
    try:
        t0 = time.time()
        resp = requests.post(OLLAMA_URL, json=payload, timeout=60)
        elapsed = time.time() - t0
        if resp.status_code != 200:
            return {"error": f"HTTP {resp.status_code}"}
        raw = resp.json().get("response", "")
        # Extract JSON from response (strip thinking tags if present)
        if "<think>" in raw:
            raw = raw[raw.rfind("</think>") + 8:].strip()
        # Find JSON block
        start = raw.find("{")
        end = raw.rfind("}") + 1
        if start == -1 or end == 0:
            return {"error": "no JSON in response", "raw": raw[:200]}
        result = json.loads(raw[start:end])
        result["_elapsed_s"] = round(elapsed, 2)
        return result
    except json.JSONDecodeError as e:
        return {"error": f"JSON parse error: {e}", "raw": raw[:200]}
    except Exception as e:
        return {"error": str(e)}


# ── Main Test ─────────────────────────────────────────────────────────────────

def main():
    print(f"Loading dataset...")
    df = pd.read_parquet(
        "data/parquet/midlarge-liquid-long-v1_bucketfix_full_audit_tier3/train.parquet",
        columns=["event_id", "ticker", "event_date", "event_type", "event_direction",
                 "parse_confidence_overall", "guidance_status", "fwd_return_5d",
                 "document_quality_score", "signal_strength_score"]
    )
    print(f"Total events: {len(df)}")

    # Build sample: mix of confidence levels and outcomes
    # Group 1: High confidence, bullish (LLM should agree with rule)
    g1 = df[
        (df["parse_confidence_overall"] >= 0.75) &
        (df["event_direction"] == "bullish") &
        (df["event_type"] == "earnings_release") &
        (df["fwd_return_5d"] > 0.03)
    ].sample(n=15, random_state=42)

    # Group 2: Low confidence, currently filtered (conf 0.40-0.60), positive outcome
    g2 = df[
        (df["parse_confidence_overall"] >= 0.40) &
        (df["parse_confidence_overall"] < 0.60) &
        (df["fwd_return_5d"] > 0.03)
    ].sample(n=15, random_state=42)

    # Group 3: Low confidence, negative outcome (these should stay filtered)
    g3 = df[
        (df["parse_confidence_overall"] >= 0.40) &
        (df["parse_confidence_overall"] < 0.60) &
        (df["fwd_return_5d"] < -0.02)
    ].sample(n=10, random_state=42)

    # Group 4: Guidance raised cases (test guidance detection)
    g4 = df[
        (df["guidance_status"] == "raised") &
        (df["parse_confidence_overall"] >= 0.65)
    ].sample(n=10, random_state=42)

    sample = pd.concat([g1, g2, g3, g4]).drop_duplicates(subset="event_id")
    sample["group"] = (
        ["high_conf_bullish"] * len(g1) +
        ["low_conf_positive"] * len(g2) +
        ["low_conf_negative"] * len(g3) +
        ["guidance_raised"] * len(g4)
    )[:len(sample)]
    print(f"Sample size: {len(sample)} events\n")

    results = []
    for i, row in enumerate(sample.itertuples()):
        accession_no = extract_accession_no(row.event_id)
        if not accession_no:
            continue

        text = get_exhibit_text(accession_no)
        if not text:
            print(f"  [{i+1}/{len(sample)}] {row.ticker} {row.event_date} — no exhibit, skip")
            continue

        # Rule parser
        rule_result = run_rule_parser(accession_no)
        # LLM parser
        llm_result = run_llm_parser(text)

        # Agreement checks
        direction_agree = (
            rule_result.get("event_direction") == llm_result.get("event_direction")
            if "error" not in llm_result else None
        )
        guidance_agree = (
            rule_result.get("guidance_status") == llm_result.get("guidance_status")
            if "error" not in llm_result else None
        )
        llm_has_eps = (
            llm_result.get("eps_beat_pct") is not None
            if "error" not in llm_result else False
        )

        record = {
            "ticker": row.ticker,
            "date": str(row.event_date),
            "accession": accession_no,
            "true_event_type": row.event_type,
            "true_direction": row.event_direction,
            "true_guidance": row.guidance_status,
            "rule_conf": round(row.parse_confidence_overall, 3),
            "fwd5d": round(row.fwd_return_5d * 100, 2),
            # Rule parser
            "rule_direction": rule_result.get("event_direction"),
            "rule_guidance": rule_result.get("guidance_status"),
            # LLM parser
            "llm_direction": llm_result.get("event_direction"),
            "llm_guidance": llm_result.get("guidance_status"),
            "llm_conf": llm_result.get("confidence"),
            "llm_eps_beat": llm_result.get("eps_beat_pct"),
            "llm_rev_beat": llm_result.get("revenue_beat_pct"),
            "llm_signal": llm_result.get("key_signal", ""),
            "llm_elapsed": llm_result.get("_elapsed_s"),
            "llm_error": llm_result.get("error"),
            # Agreement
            "direction_agree": direction_agree,
            "guidance_agree": guidance_agree,
            "llm_has_eps": llm_has_eps,
            "group": getattr(row, "group", "?"),
        }
        results.append(record)

        agree_str = "✓" if direction_agree else ("✗" if direction_agree is False else "?")
        eps_str = f" EPS:{llm_result.get('eps_beat_pct'):+.1f}%" if llm_has_eps else ""
        err_str = f" ERR:{llm_result.get('error','')[:40]}" if llm_result.get("error") else ""
        print(f"  [{i+1:2d}/{len(sample)}] {row.ticker:6s} {row.event_date} "
              f"rule_conf={row.parse_confidence_overall:.2f} fwd5d={row.fwd_return_5d*100:+.1f}%"
              f" | dir: {rule_result.get('event_direction','?'):8s} → {llm_result.get('event_direction','?'):8s} {agree_str}"
              f" | guid: {rule_result.get('guidance_status','?'):10s} → {llm_result.get('guidance_status','?'):10s}"
              f"{eps_str}{err_str}")

    if not results:
        print("No results!")
        return

    # ── Summary ───────────────────────────────────────────────────────────────
    rdf = pd.DataFrame(results)
    print("\n" + "="*80)
    print("SUMMARY")
    print("="*80)

    valid = rdf[rdf["llm_error"].isna()]
    print(f"\nProcessed: {len(results)}, LLM errors: {len(rdf[rdf['llm_error'].notna()])}")
    print(f"Avg LLM latency: {valid['llm_elapsed'].mean():.1f}s")

    print(f"\n── Direction Agreement ──")
    print(f"  Overall agreement: {valid['direction_agree'].mean()*100:.1f}%")
    for grp in valid['group'].unique():
        sub = valid[valid['group'] == grp]
        agree_pct = sub['direction_agree'].mean() * 100
        print(f"  {grp:25s}: {agree_pct:.1f}% ({len(sub)} events)")

    print(f"\n── Guidance Agreement ──")
    print(f"  Overall: {valid['guidance_agree'].mean()*100:.1f}%")

    print(f"\n── EPS Extraction (LLM only) ──")
    eps_rows = valid[valid["llm_has_eps"]]
    print(f"  EPS extracted in {len(eps_rows)}/{len(valid)} events ({len(eps_rows)/len(valid)*100:.1f}%)")
    if len(eps_rows) > 0:
        beats = eps_rows[eps_rows["llm_eps_beat"] > 0]
        misses = eps_rows[eps_rows["llm_eps_beat"] < 0]
        print(f"  Beats: {len(beats)}, Misses: {len(misses)}")
        print(f"  Avg beat %: {eps_rows['llm_eps_beat'].mean():+.2f}%")

    print(f"\n── Low-Confidence Case Analysis (rule_conf < 0.60) ──")
    low_conf = valid[valid["rule_conf"] < 0.60]
    if len(low_conf) > 0:
        print(f"  Events currently filtered (conf<0.60): {len(low_conf)}")
        print(f"  LLM avg confidence on these: {low_conf['llm_conf'].mean():.2f}")
        would_rescue = low_conf[low_conf["llm_conf"] >= 0.70]
        print(f"  LLM would rescue (LLM conf ≥ 0.70): {len(would_rescue)}")
        if len(would_rescue) > 0:
            correct = would_rescue[
                (would_rescue["llm_direction"] == "bullish") &
                (would_rescue["fwd5d"] > 0)
            ]
            print(f"  Of rescued: correct direction + positive outcome: {len(correct)}/{len(would_rescue)}")

    print(f"\n── Direction Disagreements ──")
    disagree = valid[valid["direction_agree"] == False]
    if len(disagree) > 0:
        for _, row in disagree.iterrows():
            rule_dir = row['rule_direction'] or "?"
            llm_dir = row['llm_direction'] or "?"
            print(f"  {row['ticker']:6s} {row['date']} rule={rule_dir:8s} "
                  f"llm={llm_dir:8s} fwd5d={row['fwd5d']:+.1f}% "
                  f"[{row['group']}]")
            if row["llm_signal"]:
                print(f"    LLM says: {row['llm_signal'][:100]}")

    # Save full results
    out_path = "data/analysis/llm_parser_comparison.csv"
    os.makedirs("data/analysis", exist_ok=True)
    rdf.to_csv(out_path, index=False)
    print(f"\nFull results saved to: {out_path}")


if __name__ == "__main__":
    main()