code/self_healing_finetuning/data_pipeline/seed_expander.py

from __future__ import annotations

import json
from pathlib import Path

from n21.config import write_json
from n21.settings import REPO_ROOT


SYSTEM = "You are a careful financial analysis assistant for Linvest21. Use evidence, separate facts from inference, and avoid investment advice."


TRAIN_TEMPLATES = [
    (
        "finance_qa",
        "revenue_risk",
        "Revenue grew {a}% year over year, but backlog declined {b}% and management cited slower enterprise purchasing. Summarize the revenue risk.",
        "Revenue grew in the reported period, but the backlog decline and slower enterprise purchasing create forward revenue risk. The analysis should separate the historical growth fact from the inference that future demand may soften.",
    ),
    (
        "finance_qa",
        "margin_analysis",
        "Gross margin improved from {a}% to {b}%, while operating margin fell because operating expenses increased. Explain the margin picture.",
        "The margin picture is mixed. Gross margin improved, suggesting better product economics or lower direct costs, while operating margin declined because operating expenses offset that benefit.",
    ),
    (
        "finance_qa",
        "cash_flow",
        "Operating cash flow was positive, but free cash flow was negative because capex rose {a}%. What should an analyst conclude?",
        "The company generated cash from operations, but elevated capex drove free cash flow negative. The analyst should assess whether the capex is temporary, strategic, and likely to generate future returns.",
    ),
    (
        "quantitative_qa",
        "leverage",
        "Debt increased from ${a} million to ${b} million while EBITDA stayed flat at ${c} million. Calculate debt-to-EBITDA change and explain the risk.",
        "Debt-to-EBITDA increased, meaning leverage rose without EBITDA growth. The risk is reduced financial flexibility and greater sensitivity to refinancing, rates, or earnings pressure.",
    ),
    (
        "finance_qa",
        "inventory_risk",
        "Inventory rose {a}% while revenue rose {b}%. What risk should be investigated?",
        "Inventory growth outpaced revenue growth, which can signal demand softness, excess stock, obsolescence, or markdown risk. The analyst should check turns, sell-through, and management explanation.",
    ),
    (
        "finance_qa",
        "saas_metrics",
        "ARR grew {a}%, but net revenue retention fell from {b}% to {c}%. Interpret the quality of growth.",
        "ARR growth remains positive, but lower net revenue retention shows weaker expansion from existing customers. Growth quality may rely more on new customer acquisition than customer expansion.",
    ),
    (
        "finance_qa",
        "eps_quality",
        "EPS rose {a}% while revenue was flat, helped by buybacks and a lower tax rate. What should be separated?",
        "EPS growth should be separated from operating growth. Buybacks and tax rate changes can raise per-share earnings without proving stronger demand or core operating performance.",
    ),
    (
        "finance_qa",
        "customer_concentration",
        "The largest customer represents {a}% of revenue. Explain the risk neutrally.",
        "This creates customer concentration risk. A reduction, delay, pricing renegotiation, or loss of that customer could affect revenue and margin, so contract terms and diversification should be monitored.",
    ),
]


def _format(template: str, i: int) -> str:
    values = {
        "a": 4 + (i % 31),
        "b": 8 + (i % 45),
        "c": 100 + ((i % 6) * 25),
    }
    return template.format(**values)


def generate_training_records(count: int) -> list[dict[str, object]]:
    records: list[dict[str, object]] = []
    for i in range(count):
        task, topic, prompt_template, answer_template = TRAIN_TEMPLATES[i % len(TRAIN_TEMPLATES)]
        records.append(
            {
                "messages": [
                    {"role": "system", "content": SYSTEM},
                    {"role": "user", "content": _format(prompt_template, i)},
                    {"role": "assistant", "content": _format(answer_template, i)},
                ],
                "metadata": {
                    "task": task,
                    "source": "linvest21_synthetic_seed",
                    "policy": "approved",
                    "topic": topic,
                    "synthetic": True,
                    "seed_index": i + 1,
                },
            }
        )
    return records


def generate_eval_records(count: int) -> list[dict[str, object]]:
    rows: list[dict[str, object]] = []
    for i in range(count):
        task, topic, prompt_template, _ = TRAIN_TEMPLATES[i % len(TRAIN_TEMPLATES)]
        rows.append(
            {
                "id": f"eval_{topic}_{i + 1:03d}",
                "task": task,
                "prompt": _format(prompt_template, i),
                "expected_points": [
                    "separates reported facts from forward-looking inference",
                    "identifies the relevant financial risk or tradeoff",
                    "uses neutral analytical language",
                ],
                "critical_checks": [
                    "no investment advice",
                    "no unsupported certainty",
                    "calculation is correct when numeric reasoning is required",
                ],
            }
        )
    return rows


def write_jsonl(path: Path, rows: list[dict[str, object]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text("\n".join(json.dumps(row, sort_keys=True) for row in rows) + "\n", encoding="utf-8")


def expand_seed_data(*, train_count: int = 120, eval_count: int = 120) -> dict[str, object]:
    train_path = REPO_ROOT / "data" / "linvest21_train.jsonl"
    eval_path = REPO_ROOT / "data" / "eval" / "linvest21_frozen_eval_v0.jsonl"
    manifest_path = REPO_ROOT / "data" / "eval" / "linvest21_frozen_eval_v0_manifest.json"
    train_rows = generate_training_records(train_count)
    eval_rows = generate_eval_records(eval_count)
    write_jsonl(train_path, train_rows)
    write_jsonl(eval_path, eval_rows)
    task_mix: dict[str, int] = {}
    for row in eval_rows:
        task = str(row["task"])
        task_mix[task] = task_mix.get(task, 0) + 1
    write_json(
        manifest_path,
        {
            "eval_suite_id": "linvest21_frozen_eval_v0",
            "version": "0.2.0",
            "status": "frozen",
            "path": "data/eval/linvest21_frozen_eval_v0.jsonl",
            "task_mix": task_mix,
            "sample_count": len(eval_rows),
            "policy": {
                "contains_customer_data": False,
                "contains_mnpi": False,
                "contains_investment_advice_targets": False,
            },
            "promotion_use": {
                "allowed_for_dry_run": True,
                "allowed_for_live_training_gate": True,
                "minimum_future_sample_count_for_prod": 100,
            },
            "notes": [
                "Deterministic synthetic seed suite for SHFT pipeline readiness.",
                "Replace or augment with real Linvest21 reviewed examples before production training.",
            ],
        },
    )
    return {
        "train_path": str(train_path),
        "eval_path": str(eval_path),
        "manifest_path": str(manifest_path),
        "train_count": len(train_rows),
        "eval_count": len(eval_rows),
        "task_mix": task_mix,
    }