SPB Git

spb/wp45_uqo

Public
18.4 KB · 160 lines python
1"""Phase 5 — retrieval-augmented writing engine: outline / write / upgrade a section; role-based reviews. Conditions A–E are the2experimental arms (config/evaluation.yaml); E is the default engine (profiles + structured retrieval + verification + specialist referee).3The engine never edits the manuscript: it returns text + change log + flags + audit, optionally written to a NEW output file."""4from __future__ import annotations5import json, re, time, difflib6from pathlib import Path7import yaml8from . import config, auditing, llm9from .provenance import stamp, now_iso, Ledger1011PROMPTS = config.ROOT / "prompts"12PROFILE_FOR = {("asset_pricing", "empirical"): "empirical_asset_pricing", ("corporate_finance", "empirical"): "corporate_finance", ("financial_intermediation_banking", "empirical"): "financial_intermediation_banking",13               ("household_finance", "empirical"): "household_behavioral_finance", ("behavioral_finance", "empirical"): "household_behavioral_finance", ("financial_econometrics", "*"): "financial_econometrics", ("*", "theoretical"): "theoretical_finance"}1415def load_prompt(rel: str) -> tuple[str, str, str]:16    raw = (PROMPTS / rel).read_text(); m = re.search(r'prompt_id:\s*(\S+)', raw); v = re.search(r'version:\s*"([^"]+)"', raw)17    return raw.split("---", 2)[2].strip(), (m.group(1) if m else rel), (v.group(1) if v else "0")1819def manuscript_material(ms: dict, section: str, max_chars: int = 60000, brief: str | None = None) -> str:20    """Bounded, structured description of the target manuscript handed to the model."""21    parts = [f"TITLE: {ms['title']}", f"ABSTRACT: {ms['abstract'] or '(none)'}"]22    if ms.get("macros"): parts.append("METADATA: " + json.dumps(ms["macros"]))23    if brief: parts.append("RESEARCH BRIEF (verified by the author):\n" + brief)24    parts.append("SECTIONS (heading — role — words):\n" + "\n".join(f"- {s['heading']} — {s['role']} — {s['n_words']}" for s in ms["sections"] if s["level"] == 1))25    tgt = auditing.section_text(ms, section)26    parts.append(f"CURRENT {section.upper()} SECTION:\n" + (tgt.strip()[:12000] if tgt else "(absent — to be written)"))27    others = [s for s in ms["sections"] if s["level"] == 1 and s["role"] not in (section, "appendix", "other") and s["heading"] != "(preamble)"]28    budget = max_chars - sum(len(p) for p in parts)29    for s in others:30        if budget <= 2000: break31        chunk = s["full_text"].strip()[: min(6000, budget // max(len(others), 1))]; parts.append(f"SECTION '{s['heading']}' ({s['role']}):\n{chunk}"); budget -= len(chunk)32    if ms["tables"]: parts.append("TABLES (caption | label | numbers):\n" + "\n".join(f"- {t['caption'][:200] if t['caption'] else '(no caption)'} | {t['label']} | {t['n_numbers']} numbers" for t in ms["tables"][:20]))33    if ms["equations"]: parts.append("EQUATIONS:\n" + "\n".join(f"- {e['label']}: {e['body'][:160]}" for e in ms["equations"][:20]))34    parts.append("BIBTEX KEYS AVAILABLE: " + (", ".join(ms["bib_keys"][:300]) or "(none)"))35    return "\n\n".join(parts)3637def detect_genre(ms: dict) -> dict:38    """Cheap heuristic genre/subfield detection from title+abstract+headings (no API call)."""39    t = ((ms["title"] or "") + " " + (ms["abstract"] or "") + " " + " ".join(s["heading"] for s in ms["sections"])).lower()40    genre = "theoretical" if re.search(r"\b(theorem|proposition|lemma|proof|equilibrium model)\b", t) and not re.search(r"\b(regression|sample|data set|dataset|panel|estimates?)\b", t) else "empirical"41    subs = {"asset_pricing": r"asset pricing|return predictab|risk premi|factor|cross-section of|volatility|options?|bond yields?|forecast", "corporate_finance": r"corporate|firm|investment|capital structure|leverage|m&a|merger|governance|ceo|payout|dividend",42            "financial_intermediation_banking": r"bank|lender|loan|credit|intermediar|deposit", "household_finance": r"household|consumer|mortgage|retirement|investor behavior|individual investor", "behavioral_finance": r"behavioral|bias|sentiment|attention|overconfiden",43            "market_microstructure": r"microstructure|order flow|liquidity|high-frequency|limit order|market maker|intraday", "financial_econometrics": r"estimator|econometric|inference|monte carlo|asymptotic|realized volatility|forecasting architecture", "macro_finance": r"monetary policy|macro|business cycle|sovereign|exchange rate",44            "real_estate_housing": r"housing|real estate|hedonic|property|rent", "ai_agents_methods": r"language model|llm|agent|claude|gpt|autonomous|machine learning"}45    scores = {k: len(re.findall(v, t)) for k, v in subs.items()}; sub = max(scores, key=scores.get) if max(scores.values()) > 0 else "other"46    return {"genre": genre, "subfield": sub, "scores": scores}4748def profile_for(genre: str, subfield: str, override: str | None = None) -> tuple[str, dict, str]:49    name = override or PROFILE_FOR.get((subfield, genre)) or PROFILE_FOR.get((subfield, "*")) or PROFILE_FOR.get(("*", genre)) or ("empirical_all" if genre == "empirical" else "corpus_all")50    p = config.ROOT / "profiles" / f"{name}.yaml"51    if not p.exists(): name, p = "corpus_all", config.ROOT / "profiles" / "corpus_all.yaml"52    prof = yaml.safe_load(p.read_text()) if p.exists() else {}53    if prof.get("n_documents", 0) < 8 and name != "corpus_all":54        fb = "empirical_all" if genre == "empirical" else "corpus_all"; p2 = config.ROOT / "profiles" / f"{fb}.yaml"55        if p2.exists() and yaml.safe_load(p2.read_text()).get("n_documents", 0) > prof.get("n_documents", 0): name, prof = fb, yaml.safe_load(p2.read_text())56    note = ""57    if subfield in ("real_estate_housing", "ai_agents_methods", "other") or genre == "theoretical":58        note = ("Note: this manuscript is outside the corpus's core empirical-finance genre; JoF conventions below apply only where the argument structure is comparable "59                "(question → contribution → evidence → bounded claims). Do not force an empirical-paper template on a theory, methods or computer-science manuscript; explain in the change log which conventions you did not apply and why.")60    return name, prof, note6162def profile_summary(prof: dict, section: str) -> str:63    if not prof: return "(no profile available)"64    st = prof.get("structure", {}); rh = prof.get("rhetoric", {}); L = [f"n = {prof.get('n_documents')} documents; {prof.get('caveat')}"]65    def md(k, src=st): v = src.get(k, {}); return f"median {v.get('median')} (IQR {v.get('q25')}–{v.get('q75')}, n={v.get('n')})" if v.get("n") else "n/a"66    L.append(f"Body length: {md('n_words_body')} words; introduction {md('intro_words')} words, {md('intro_paragraphs')} paragraphs, share of body {md('intro_share')}; conclusion {md('conclusion_words')} words; tables {md('n_tables')}; figures {md('n_figures')}; numbered equations {md('n_equations')}; references {md('n_references')}.")67    if rh:68        am = rh.get("abstract_moves_present", {}); L.append("Abstract moves present (share): " + ", ".join(f"{k} {v.get('share')}" for k, v in am.items()) + f"; abstract sentences {md('abstract_sentences', rh)}.")69        intro = rh.get("introduction", {}); L.append(f"Introduction: first move {intro.get('first_move')}; common move sequences {intro.get('common_move_sequences')}; contribution paragraph relative position {md('contribution_relative_position', intro)}; roadmap paragraph present (LLM) {intro.get('roadmap_present_llm', {}).get('share')}; literature positioning {intro.get('literature_positioning_location')}; contribution types {intro.get('contribution_types')}.")70        L.append(f"Identification: causal-language strength {md('causal_language_strength', rh.get('identification', {}))}/5; results: economic significance discussed {rh.get('results', {}).get('economic_significance_discussed', {}).get('share')}, statistical language {rh.get('results', {}).get('statistical_language')}; robustness location {rh.get('robustness', {}).get('location')}; table SE conventions {rh.get('tables', {}).get('standard_errors')}; conclusion bounded claims {rh.get('conclusion', {}).get('bounded_claims', {}).get('share')}, open questions {rh.get('conclusion', {}).get('open_questions', {}).get('share')}; hedging {rh.get('hedging')}.")71    return "\n".join(L)7273def retrieve_for(ms: dict, section: str, unit: str, k: int = 6, genre: str | None = None, subfield: str | None = None) -> list[dict]:74    from . import retrieval75    q = f"{ms['title']}. {(ms['abstract'] or '')[:600]}"; role = {"introduction": "introduction", "literature": "literature", "conclusion": "conclusion", "data": "data", "results": "results"}.get(section)76    hits = retrieval.search(q, k=k, unit=unit, section=role if unit == "raw" else None, subfield=subfield if subfield in ("asset_pricing", "corporate_finance", "financial_intermediation_banking", "household_finance", "behavioral_finance", "market_microstructure", "financial_econometrics", "macro_finance") else None)77    if len(hits) < k // 2: hits = retrieval.search(q, k=k, unit=unit)78    return hits7980def fmt_hits(hits: list[dict]) -> str:81    return "\n".join(f"[{i+1}] ({h['unit']}; {h['provenance']['doc_id']} {h['provenance']['year']}, p.{h['provenance']['page']}, {h['provenance']['role']}) {h['text']}" for i, h in enumerate(hits)) or "(no retrieval)"8283def parse_output(text: str) -> dict:84    body, log, flags = text, "", ""85    if "=== CHANGE LOG ===" in text:86        body, rest = text.split("=== CHANGE LOG ===", 1)87        if "=== FLAGS ===" in rest: log, flags = rest.split("=== FLAGS ===", 1)88        else: log = rest89    body = re.sub(r"^```(?:latex|tex)?\s*|\s*```$", "", body.strip(), flags=re.S)90    return {"text": body.strip(), "change_log": log.strip(), "flags": flags.strip()}9192def run(paper: Path, section: str = "introduction", mode: str = "write", output: Path | None = None, profile: str | None = None, condition: str = "E_full_engine",93        brief: str | None = None, model: str | None = None, effort: str | None = None, ablation: dict | None = None, meta: dict | None = None) -> dict:94    """mode: outline | write | upgrade. condition: A_baseline | B_strong_prompt | C_conventional_rag | D_profiles_retrieval | E_full_engine.95    ablation (E only): {"no_retrieval", "no_profile", "no_referee", "no_factcheck"} flags."""96    t0 = time.time(); cfg = config.load("models"); model = model or cfg["writing_model"]; effort = effort or cfg["effort_writing"]; ab = ablation or {}97    ms = auditing.load_manuscript(paper); det = detect_genre(ms); cond = condition[0]98    sysb, sid, sv = load_prompt("writing/system_base_v1.md"); material = manuscript_material(ms, section, brief=brief)99    task_verb = {"outline": f"Produce a detailed OUTLINE (paragraph-by-paragraph plan with the rhetorical move of each paragraph and the evidence it will cite) for the {section}", "write": f"Write the {section}", "upgrade": f"Rewrite and upgrade the existing {section}, preserving every number, citation, claim and result"}[mode]100    retrieved = []; prof_name = None; prof = {}101    if cond == "A": body, pid, pv = load_prompt("writing/condition_A_v1.md"); user_instr = body.replace("Write the {section}", task_verb)102    elif cond == "B": body, pid, pv = load_prompt("writing/condition_B_v1.md"); user_instr = body.replace("Write the {section}", task_verb)103    elif cond == "C":104        body, pid, pv = load_prompt("writing/condition_C_v1.md"); retrieved = retrieve_for(ms, section, "raw", k=6)105        user_instr = body.replace("Write the {section}", task_verb).replace("{retrieved}", fmt_hits(retrieved))106    else:  # D or E107        body, pid, pv = load_prompt("writing/condition_D_v1.md")108        prof_name, prof, gnote = profile_for(det["genre"], det["subfield"], profile)109        if ab.get("no_profile"): prof_name, prof = "(ablated)", {}110        retrieved = [] if ab.get("no_retrieval") else retrieve_for(ms, section, "annotation", k=6, genre=det["genre"], subfield=det["subfield"])111        user_instr = (body.replace("Write the {section}", task_verb).replace("{genre}", det["genre"]).replace("{subfield}", det["subfield"]).replace("{profile_name}", str(prof_name)).replace("{profile_n}", str(prof.get("n_documents", 0)))112                      .replace("{profile_summary}", profile_summary(prof, section) if prof else "(profile ablated)").replace("{retrieved}", fmt_hits(retrieved)).replace("{genre_note}", gnote))113    user = material + "\n\n=== TASK ===\n" + user_instr114    r1 = llm.call(sysb, user, model=model, max_tokens=cfg["max_tokens_writing"], effort=effort, purpose=f"write:{condition}:{mode}", prompt_id=pid, prompt_version=pv, use_cache=False, meta={"section": section, "paper": str(paper), **(meta or {})})115    out = parse_output(r1["text"]); calls = [r1]; audit = None; referee = None116    src_text = material + "\n" + (brief or "")117    if cond == "E":118        if not ab.get("no_factcheck"):119            audit = {"numeric": auditing.numeric_preservation(src_text, out["text"]), "citations": auditing.citation_check(out["text"], ms["bib_keys"]), "unsupported": auditing.unsupported_claims(out["text"]), "overlap": auditing.overlap_with_corpus(out["text"])}120        if not ab.get("no_referee"):121            eb, eid, ev = load_prompt("writing/condition_E_referee_v1.md")122            audit_txt = json.dumps({k: (v if k != "overlap" else {"max_overlap_share": v["max_overlap_share"]}) for k, v in (audit or {}).items()}, ensure_ascii=False)[:6000] if audit else "(fact-check ablated: no deterministic audit available)"123            u2 = material + f"\n\n=== DRAFT {section.upper()} ===\n" + out["text"] + "\n\n=== DETERMINISTIC AUDIT ===\n" + audit_txt + "\n\n=== TASK ===\n" + eb.replace("{section}", section).replace("{retrieved}", "")124            r2 = llm.call(sysb, u2, model=model, max_tokens=cfg["max_tokens_writing"], effort=effort, purpose=f"referee:{condition}:{mode}", prompt_id=eid, prompt_version=ev, use_cache=False, meta={"section": section, "paper": str(paper), **(meta or {})})125            referee = parse_output(r2["text"]); calls.append(r2); draft_v1 = out; out = referee; out["change_log"] = draft_v1["change_log"] + "\n--- referee pass ---\n" + referee["change_log"]126        if not ab.get("no_factcheck"):127            audit["final"] = {"numeric": auditing.numeric_preservation(src_text, out["text"]), "citations": auditing.citation_check(out["text"], ms["bib_keys"]), "unsupported": auditing.unsupported_claims(out["text"]), "overlap": auditing.overlap_with_corpus(out["text"])}128    res = {"paper": str(paper), "section": section, "mode": mode, "condition": condition, "ablation": ab, "model": model, "effort": effort, "detected": det, "profile": prof_name, "retrieved": [{"unit": h["unit"], "score": h["score"], **h["provenance"]} for h in retrieved],129           "text": out["text"], "change_log": out["change_log"], "flags": out["flags"], "audit": audit, "usage": {"input": sum(c["usage"]["input"] + c["usage"]["cache_read"] + c["usage"]["cache_write"] for c in calls), "output": sum(c["usage"]["output"] for c in calls)},130           "cost_usd": round(sum(c["cost_usd"] for c in calls), 5), "latency_s": round(time.time() - t0, 1), "n_calls": len(calls), "prompt_versions": {pid: pv}, "generated": stamp(), "manuscript_sha256": ms["sha256"]}131    if output:132        output = Path(output); output.parent.mkdir(parents=True, exist_ok=True)133        existing = auditing.section_text(ms, section)134        hdr = f"% jof250 {mode} — section: {section}; condition {condition}; model {model}; {res['generated']['timestamp']}; manuscript sha256 {ms['sha256'][:12]}\n% This file is a PROPOSAL. The manuscript was not modified.\n"135        output.write_text(hdr + res["text"] + "\n")136        side = output.with_suffix(".report.md")137        diff = "\n".join(difflib.unified_diff((existing or "").splitlines(), res["text"].splitlines(), "current", "proposed", lineterm="")) if existing else "(section did not exist)"138        side.write_text(f"# jof250 {mode} report — {section}\n\nCondition: {condition} · model {model} · effort {effort} · profile {prof_name} · cost {res['cost_usd']} USD · latency {res['latency_s']} s\n\n## Change log\n\n{res['change_log']}\n\n## Flags\n\n{res['flags']}\n\n## Retrieval provenance\n\n" + "\n".join(f"- {h['unit']} — {h['doc_id']} ({h['year']}, p.{h['page']}, {h['role']}) score {h['score']}" for h in res["retrieved"]) + f"\n\n## Deterministic audit\n\n```json\n{json.dumps(audit, indent=1, ensure_ascii=False)[:6000] if audit else 'n/a'}\n```\n\n## Diff vs current section\n\n```diff\n{diff[:20000]}\n```\n")139        res["output"] = str(output); res["report"] = str(side)140    Ledger(config.ROOT / "data" / "annotations" / "engine_runs.jsonl").append({k: v for k, v in res.items() if k not in ("text", "change_log", "flags", "audit")})141    return res142143def review(paper: Path, output: Path | None = None, role: str = "referee", model: str | None = None, effort: str | None = None, with_conventions: bool = True) -> dict:144    t0 = time.time(); cfg = config.load("models"); model = model or cfg["writing_model"]; effort = effort or cfg["effort_writing"]145    ms = auditing.load_manuscript(paper); det = detect_genre(ms); material = manuscript_material(ms, "results")146    rep = auditing.audit_manuscript(paper, overlap=False); facts = json.dumps({k: rep[k] for k in ("tables", "equations", "citations", "unsupported_claims") if k in rep}, ensure_ascii=False)[:8000]147    body, pid, pv = load_prompt(f"review/{role}_v1.md"); conv = ""148    if with_conventions:149        name, prof, note = profile_for(det["genre"], det["subfield"]); conv = f"\nMeasured conventions of the reference corpus (profile `{name}`): {profile_summary(prof, 'all')} {note}"150    sysb = "You are a rigorous, fair reviewer. Never invent results; tie every comment to the material; keep every number unchanged."151    user = material + "\n\n=== DETERMINISTIC AUDIT FACTS ===\n" + facts + "\n\n=== TASK ===\n" + body.replace("{conventions}", conv)152    r = llm.call(sysb, user, model=model, max_tokens=cfg["max_tokens_writing"], effort=effort, purpose=f"review:{role}", prompt_id=pid, prompt_version=pv, use_cache=False, meta={"paper": str(paper)})153    res = {"paper": str(paper), "role": role, "model": model, "detected": det, "text": r["text"], "cost_usd": r["cost_usd"], "latency_s": round(time.time() - t0, 1), "generated": stamp(), "audit_facts": rep}154    if output:155        Path(output).parent.mkdir(parents=True, exist_ok=True)156        Path(output).write_text(f"# jof250 review ({role}) — {ms['title']}\n\n_Model {model}, effort {effort}, {res['generated']['timestamp']}; cost {r['cost_usd']} USD. Deterministic audit appended._\n\n" + r["text"] + "\n\n---\n\n" + auditing.render_report(rep))157        res["output"] = str(output)158    Ledger(config.ROOT / "data" / "annotations" / "engine_runs.jsonl").append({k: v for k, v in res.items() if k not in ("text", "audit_facts")})159    return res160
Files · 09f9e1493a5c5e2614d317520925ffaafec79679
Menu

Theme