research_agent.py

"""
research_agent.py — Auto Research: a multi-step agent, fully local (llama.cpp).

The agent loop (every step is logged to a JSON trace → 📡 "Sharing is Caring"):

    PLAN      LLM drafts 3-5 research questions for the ticker
    TOOLS     deterministic tool calls gather evidence:
                t_fundamentals  yfinance .info (valuation, margins, growth…)
                t_financials    last quarterly income-statement lines
                t_price         52w range, MA200, drawdown, momentum (data_us)
                t_chan          the user's Chan engine verdict (signal_runner)
                t_news          recent headlines
    ANALYZE   LLM writes each report section against the gathered evidence:
                Valuation · Moat & supply-chain position · Bull case ·
                Bear case · Technical timing (Chan) · Risks & verdict
    REPORT    sections assembled into a markdown report, saved to /data/reports
    TRACE     full step-by-step trace saved to /data/traces/<ticker>_<ts>.json
              (upload that folder as a HF dataset to claim the open-trace badge)

Feature 4 (auto-trigger) lives in automation.py: whenever a NEW ticker enters
the signal pool, this agent runs for it automatically after the daily pipeline.
"""
from __future__ import annotations

import datetime as dt
import json
import os

import pandas as pd

import paths

# Multi-agent split: the 4B Analyst writes the heavy sections while the 1.7B
# Reporter writes market/technical sections IN PARALLEL on its own lock —
# wall-clock time ≈ the slower of the two instead of their sum.
ANALYST_SECTIONS = [
    ("Valuation", "cheap/fair/rich vs the growth and margins in evidence; quote 2-3 numbers"),
    ("Technology moat", "the company's technological moat and how defensible it is"),
    ("Supply-chain map", "a markdown table with two columns 'Upstream suppliers' and "
     "'Downstream customers/users': list 4-6 real companies on each side WITH stock "
     "tickers in parentheses, e.g. TSMC (TSM); mark private companies (private)"),
    ("Bull case", "strongest 3 points FOR owning it"),
    ("Bear case", "strongest 3 points AGAINST owning it"),
]
REPORTER_SECTIONS = [
    ("Money flow & related tickers", "read the MONEY FLOW evidence: is capital "
     "entering or leaving the stock and its sector; name the sector ETF and 2-4 "
     "related tickers worth watching"),
    ("Technical timing (Chan theory)", "interpret the CHAN ENGINE VERDICT for a "
     "long-term holder: act now, wait, or exit, and the key price levels"),
    ("Risks & verdict", "top risks, then one line: Buy / Accumulate / Hold / Avoid"),
]
SECTIONS = ANALYST_SECTIONS + REPORTER_SECTIONS  # kept for trace readability


# ───────────────────────── trace plumbing ─────────────────────────
class Trace:
    def __init__(self, ticker: str):
        self.ticker = ticker
        self.t0 = dt.datetime.utcnow()
        self.steps = []

    def log(self, step: str, kind: str, content):
        self.steps.append({
            "n": len(self.steps) + 1,
            "t": dt.datetime.utcnow().isoformat(timespec="seconds") + "Z",
            "step": step, "type": kind,
            "content": content if isinstance(content, (str, dict, list)) else str(content),
        })

    def save(self) -> str:
        ts = self.t0.strftime("%Y%m%d-%H%M%S")
        path = os.path.join(paths.TRACES_DIR, f"{self.ticker}_{ts}.json")
        try:
            with open(path, "w", encoding="utf-8") as f:
                json.dump({"ticker": self.ticker, "agent": "chan-compass-research",
                           "model_runtime": "llama.cpp (local)",
                           "started": self.t0.isoformat() + "Z",
                           "steps": self.steps}, f, ensure_ascii=False, indent=1, default=str)
            return path
        except OSError:
            return ""


def _llm(prompt: str, max_tokens: int = 500) -> str:
    try:
        import llm_local
        if llm_local.is_loaded("deep"):
            return llm_local.chat(prompt, max_tokens=max_tokens, worker="deep")
        if llm_local.is_loaded("fast"):       # deep still loading → fall back
            return llm_local.chat(prompt, max_tokens=max_tokens, worker="fast")
    except Exception:
        pass
    return ""


# ───────────────────────── tools ─────────────────────────
def t_fundamentals(ticker: str) -> dict:
    import research
    md, plain, err = research.gather_facts(ticker)
    return {"ok": not err, "markdown": md, "plain": plain, "error": err}


def t_financials(ticker: str) -> str:
    try:
        import yfinance as yf
        q = yf.Ticker(ticker).quarterly_income_stmt
        if q is None or q.empty:
            return ""
        rows = [r for r in ("Total Revenue", "Gross Profit", "Operating Income",
                            "Net Income") if r in q.index]
        sub = q.loc[rows].iloc[:, :4]
        out = []
        for r in rows:
            vals = ", ".join(f"{c.strftime('%Y-%m')}: ${v/1e9:,.2f}B"
                             for c, v in sub.loc[r].items() if pd.notna(v))
            out.append(f"{r} — {vals}")
        return "\n".join(out)
    except Exception:
        return ""


def t_price(ticker: str) -> str:
    try:
        import data_us
        d = data_us.load_level(ticker, "d")
        if d is None or len(d) < 60:
            return ""
        px = float(d["close"].iloc[-1])
        hi52 = float(d["close"].tail(252).max()); lo52 = float(d["close"].tail(252).min())
        ma200 = float(d["close"].rolling(200).mean().iloc[-1])
        r3m = px / float(d["close"].iloc[-63]) - 1 if len(d) > 63 else 0
        r1y = px / float(d["close"].iloc[-252]) - 1 if len(d) > 252 else 0
        return (f"Price ${px:,.2f} | 52w range ${lo52:,.2f}–${hi52:,.2f} "
                f"({(px/hi52-1):+.1%} off high) | vs MA200 {(px/ma200-1):+.1%} | "
                f"3m {r3m:+.1%}, 1y {r1y:+.1%}")
    except Exception:
        return ""


def t_chan(ticker: str) -> str:
    try:
        import signal_runner
        row, _ = signal_runner.analyze_one(ticker)
        if not row:
            return ""
        return (f"Tomorrow: {row['Tomorrow']} | signal {row['Signal']} | "
                f"confidence {row['Confidence']} | buy zone {row['Buy zone']} | "
                f"invalid below {row['Invalid below']} | note: {row['Note']}")
    except Exception:
        return ""


def t_flows(ticker: str) -> str:
    """Money-flow proxy (Δ% × dollar volume) for the ticker and its sector ETF,
    1/5/20-day windows, plus related tickers via the sector mapping."""
    try:
        import data_us
        import rotation as rot
        out = []
        d = data_us.load_level(ticker, "d")
        for n, lab in ((1, "1D"), (5, "5D"), (20, "20D")):
            st = rot._window_stats(d, n)
            if st:
                pct, dvol, flow = st
                out.append(f"{ticker} {lab}: {pct:+.2%}, flow proxy "
                           f"${flow/1e6:+,.0f}M on ${dvol/1e9:,.1f}B avg $vol")
        sector = ""
        try:
            import yfinance as yf
            sector = (yf.Ticker(ticker).info or {}).get("sector", "")
        except Exception:
            pass
        etf = next((k for k, v in rot.SECTOR_ETFS.items() if v == sector), None)
        if etf:
            de = data_us.load_level(etf, "d")
            st = rot._window_stats(de, 5)
            if st:
                out.append(f"Sector ETF {etf} ({sector}) 5D: {st[0]:+.2%}, "
                           f"flow proxy ${st[2]/1e6:+,.0f}M")
        return "\n".join(out)
    except Exception:
        return ""


def t_news(ticker: str) -> str:
    try:
        import yfinance as yf
        heads = []
        for x in (yf.Ticker(ticker).news or [])[:8]:
            c = x.get("content", x)
            t = c.get("title")
            if t:
                heads.append("- " + t)
        return "\n".join(heads)
    except Exception:
        return ""


# ───────────────────────── the agent ─────────────────────────
def _gather_evidence(ticker: str, tr: "Trace", on_step=None) -> dict:
    """Run all evidence tools IN PARALLEL (network-bound) — was serial before."""
    from concurrent.futures import ThreadPoolExecutor
    tools = {"fundamentals": t_fundamentals, "financials": t_financials,
             "price": t_price, "chan_engine": t_chan, "flows": t_flows,
             "news": t_news}
    evidence = {}
    with ThreadPoolExecutor(max_workers=6) as ex:
        futs = {name: ex.submit(fn, ticker) for name, fn in tools.items()}
        for name, fut in futs.items():
            try:
                evidence[name] = fut.result(timeout=40)
            except Exception as e:
                evidence[name] = ""
                tr.log(f"TOOL:{name}", "tool_error", str(e))
            brief = (evidence[name].get("plain", "")[:300]
                     if isinstance(evidence[name], dict) else str(evidence[name])[:300])
            tr.log(f"TOOL:{name}", "tool_result", brief or "(empty)")
            if on_step:
                on_step(name)
    return evidence


def _evidence_text(evidence: dict) -> str:
    fund = evidence.get("fundamentals")
    return (f"FUNDAMENTALS:\n{(fund.get('plain','') if isinstance(fund, dict) else '')[:1000]}\n\n"
            f"QUARTERLY FINANCIALS:\n{str(evidence.get('financials'))[:380] or 'n/a'}\n\n"
            f"PRICE ACTION:\n{evidence.get('price') or 'n/a'}\n\n"
            f"MONEY FLOW:\n{str(evidence.get('flows'))[:420] or 'n/a'}\n\n"
            f"CHAN ENGINE VERDICT:\n{str(evidence.get('chan_engine'))[:380] or 'n/a'}\n\n"
            f"RECENT HEADLINES:\n{str(evidence.get('news'))[:380] or 'n/a'}")[:2500]


def _sections_prompt(ticker: str, sections, ev_text: str) -> str:
    sec_list = "\n".join(f"## {t} — {i}" for t, i in sections)
    return (f"You are a buy-side equity analyst. Write part of a research note on "
            f"{ticker} using ONLY the evidence below (write 'n/a' for missing facts, "
            f"never invent numbers; company names in the supply-chain table may come "
            f"from your own industry knowledge). Output EXACTLY these markdown "
            f"sections, ≤70 words each, plain prose, ENGLISH ONLY, no disclaimers:\n"
            f"{sec_list}\n\nEVIDENCE:\n{ev_text}")


_STATIC_PLAN = ("1. Is the valuation justified by growth?\n"
                "2. How durable is the technology moat and supply-chain position?\n"
                "3. Where is capital flowing — into or out of the name and sector?\n"
                "4. Is now a good technical entry for a long-term holder?")


def run_research(ticker: str, auto: bool = False) -> tuple:
    """Blocking multi-agent run (used by the daily pipeline).
    Analyst (4B) and Reporter (1.7B) write their sections in PARALLEL.
    Returns (report_markdown, trace_path). If no model is ready, returns
    ('', '') so the caller can postpone instead of saving an empty report."""
    import llm_local
    ticker = (ticker or "").strip().upper()
    if not ticker:
        return "Enter a ticker symbol first.", ""
    if not (llm_local.is_loaded("analyst") or llm_local.is_loaded("reporter")):
        return "", ""   # postpone — model still loading
    tr = Trace(ticker)
    tr.log("PLAN", "static", _STATIC_PLAN)
    evidence = _gather_evidence(ticker, tr)
    fund = evidence.get("fundamentals")
    if isinstance(fund, dict) and not fund.get("ok"):
        tr.save()
        return f"⚠️ {fund.get('error', 'Could not fetch data.')}", ""
    ev_text = _evidence_text(evidence)

    import threading
    parts = {}

    def _write(slot, sections, worker, fallback_worker):
        wk = worker if llm_local.is_loaded(worker) else fallback_worker
        tr.log(f"ANALYZE:{slot}", "llm_request", f"worker={wk}")
        txt = llm_local.chat(_sections_prompt(ticker, sections, ev_text),
                             max_tokens=620 if slot == "analyst" else 360, worker=wk)
        if txt.startswith("(") or txt.startswith("⏳"):
            txt = ""
        tr.log(f"ANALYZE:{slot}", "llm_response", txt[:500])
        parts[slot] = txt

    th = threading.Thread(target=_write,
                          args=("reporter", REPORTER_SECTIONS, "reporter", "analyst"))
    th.start()
    _write("analyst", ANALYST_SECTIONS, "analyst", "reporter")
    th.join(timeout=300)

    fund_md = fund.get("markdown", "") if isinstance(fund, dict) else ""
    stamp = dt.datetime.utcnow().strftime("%Y-%m-%d %H:%M UTC")
    head = (f"# {ticker} — Research Note{' (auto-generated)' if auto else ''}\n"
            f"_{stamp} · multi-agent: Analyst (Qwen3-4B) + Reporter (Qwen3-1.7B), "
            f"both llama.cpp local_\n\n**Agent plan:**\n{_STATIC_PLAN}\n\n{fund_md}\n\n---\n")
    body = "\n\n".join(p for p in (parts.get("analyst"), parts.get("reporter")) if p)
    if not body:
        tr.save()
        return "", ""   # model produced nothing — postpone, never save a stub
    report = head + body + "\n\n---\n_Agent trace saved — see the Automation tab._"
    tr.log("REPORT", "assembled", f"{len(report)} chars")
    trace_path = tr.save()
    try:
        rp = os.path.join(paths.REPORTS_DIR, f"{ticker}_{tr.t0.strftime('%Y%m%d')}.md")
        with open(rp, "w", encoding="utf-8") as f:
            f.write(report)
    except OSError:
        pass
    return report, trace_path


def list_reports() -> list:
    try:
        fs = [f for f in os.listdir(paths.REPORTS_DIR) if f.endswith(".md")]
        fs.sort(key=lambda f: os.path.getmtime(os.path.join(paths.REPORTS_DIR, f)),
                reverse=True)
        return fs
    except OSError:
        return []


def read_report(fname: str) -> str:
    if not fname:
        return "No report selected."
    try:
        with open(os.path.join(paths.REPORTS_DIR, os.path.basename(fname)),
                  encoding="utf-8") as f:
            return f.read()
    except OSError as e:
        return f"Could not read report: {e}"


def list_traces() -> str:
    try:
        fs = [f for f in os.listdir(paths.TRACES_DIR) if f.endswith(".json")]
        fs.sort(key=lambda f: os.path.getmtime(os.path.join(paths.TRACES_DIR, f)),
                reverse=True)
        fs = fs[:20]
        if not fs:
            return "_No traces yet — run a research note first._"
        return ("**Saved agent traces** (`" + paths.TRACES_DIR + "`):\n" +
                "\n".join(f"- `{f}`" for f in fs) +
                "\n\n📡 To claim the open-trace badge: download this folder and "
                "push it to the Hub as a dataset (`huggingface-cli upload`).")
    except OSError:
        return "_Trace folder unavailable._"


# ───────────────────────── streaming UI runner ─────────────────────────
def run_research_stream(ticker: str):
    """Generator for the UI. Multi-agent: evidence tools run in parallel; the
    1.7B Reporter writes its sections in a background thread while the 4B
    Analyst STREAMS its sections live; never saves a 'model busy' stub."""
    import threading
    import llm_local
    ticker = (ticker or "").strip().upper()
    if not ticker:
        yield "Enter a ticker symbol first.", ""
        return
    tr = Trace(ticker)
    log_lines = [f"### 🤖 Multi-agent research · {ticker}"]

    def show(msg):
        log_lines.append(f"- {msg}")
        return "\n".join(log_lines)

    tr.log("PLAN", "static", _STATIC_PLAN)
    yield show("**PLAN** ready ✓ · **TOOLS** — gathering 6 evidence sources in parallel…"), ""
    evidence = _gather_evidence(ticker, tr)
    yield show("Evidence in: fundamentals · financials · price · money flow · Chan engine · news ✓"), ""

    fund = evidence.get("fundamentals")
    if isinstance(fund, dict) and not fund.get("ok"):
        tr.save()
        yield show(f"⚠️ {fund.get('error', 'Data fetch failed.')}"), ""
        return
    ev_text = _evidence_text(evidence)
    fund_md = fund.get("markdown", "") if isinstance(fund, dict) else ""
    stamp = dt.datetime.utcnow().strftime("%Y-%m-%d %H:%M UTC")
    head = (f"# {ticker} — Research Note\n_{stamp} · multi-agent: Analyst (Qwen3-4B) "
            f"+ Reporter (Qwen3-1.7B), both llama.cpp local_\n\n"
            f"**Agent plan:**\n{_STATIC_PLAN}\n\n{fund_md}\n\n---\n")

    # Reporter sub-agent works in parallel on its own lock
    side = {"txt": ""}

    def _side():
        wk = "reporter" if llm_local.is_loaded("reporter") else "analyst"
        t = llm_local.chat(_sections_prompt(ticker, REPORTER_SECTIONS, ev_text),
                           max_tokens=360, worker=wk)
        side["txt"] = "" if (t.startswith("(") or t.startswith("⏳")) else t
        tr.log("ANALYZE:reporter", "llm_response", side["txt"][:400])

    th = None
    if llm_local.is_loaded("reporter") or llm_local.is_loaded("analyst"):
        th = threading.Thread(target=_side, daemon=True)
        th.start()
        yield show("**Reporter sub-agent** writing money-flow / Chan timing / verdict "
                   "in parallel…"), head
    # Analyst streams the main sections live
    main = ""
    wk_main = "analyst" if llm_local.is_loaded("analyst") else "reporter"
    if llm_local.is_loaded(wk_main):
        yield show("**Analyst sub-agent** streaming valuation / moat / supply-chain "
                   "map / bull-bear…"), head
        for acc in llm_local.chat_stream(
                _sections_prompt(ticker, ANALYST_SECTIONS, ev_text),
                max_tokens=620, worker=wk_main):
            if acc.startswith("⏳") or acc.startswith("("):
                continue
            main = acc
            yield "\n".join(log_lines), head + main
    tr.log("ANALYZE:analyst", "llm_response", main[:500])
    if th is not None:
        th.join(timeout=240)
    body = "\n\n".join(p for p in (main, side["txt"]) if p)
    if not body:
        tr.save()
        yield show("⚠️ Sub-agents not ready yet (still loading) — evidence gathered "
                   "above; try again in a minute. Nothing was saved."), head
        return
    report = head + body + "\n\n---\n_Agent trace saved — see the Automation tab._"
    trace_path = tr.save()
    try:
        rp = os.path.join(paths.REPORTS_DIR, f"{ticker}_{tr.t0.strftime('%Y%m%d')}.md")
        with open(rp, "w", encoding="utf-8") as f:
            f.write(report)
    except OSError:
        pass
    yield show(f"**DONE** ✓ report + trace saved (`{os.path.basename(trace_path)}`)"), report