Manager: fixes from PR #36 review

agents/aurora_processing.py

@@ -72,10 +72,11 @@ def _assign_label(pct_change, threshold):
 
 
 def _assign_split(df, dataset_end=None):
-    """Add a time-based train/test `split` column: earliest 80% of dates =
-    `train`, the rest = `test`. Date-based (not random) so no future information
-    can leak into training — the classifier trains on `train` rows and is scored
-    only on `test` rows it has never seen.
+    """Add a time-based train/val/test `split` column: earliest 80% of dates =
+    training period, the rest = `test`. The LAST 10% of training dates become
+    `val` — the retune loop scores itself on those rows, keeping `test`
+    untouched until the final report (tuning on test would overfit it).
+    Date-based (not random) so no future information can leak into training.
 
     `dataset_end` (YYYY-MM-DD) optionally drops later rows *before* splitting —
     e.g. "2019-12-31" keeps the COVID crash out of the test window so the
@@ -85,12 +86,16 @@ def _assign_split(df, dataset_end=None):
         before = len(df)
         df = df[pd.to_datetime(df["date"]) <= pd.to_datetime(dataset_end)].reset_index(drop=True)
         print(f"  Dataset ended at {dataset_end}: dropped {before - len(df):,} later rows")
-    split_cut = pd.to_datetime(df["date"]).quantile(0.8)
-    df["split"] = ["train" if d <= split_cut else "test"
-                   for d in pd.to_datetime(df["date"])]
+    dates = pd.to_datetime(df["date"])
+    split_cut = dates.quantile(0.8)
+    val_cut = dates[dates <= split_cut].quantile(0.9)  # last 10% of training dates
+    df["split"] = ["test" if d > split_cut else ("val" if d > val_cut else "train")
+                   for d in dates]
     n_train = int((df["split"] == "train").sum())
+    n_val = int((df["split"] == "val").sum())
     n_test = int((df["split"] == "test").sum())
-    print(f"  Split at {split_cut.date()}: {n_train:,} train / {n_test:,} test")
+    print(f"  Split at {split_cut.date()} (val from {val_cut.date()}): "
+          f"{n_train:,} train / {n_val:,} val / {n_test:,} test")
     return df
 
 

agents/contracts.py

@@ -0,0 +1,144 @@
+"""Shared data-contract helpers for the agent handoff files.
+
+The source of truth is docs/data_contracts.md. This module keeps the repeated
+column lists, label validation, and common row shaping behind one interface so
+agent modules do not each reimplement the same contract details.
+"""
+
+from __future__ import annotations
+
+import csv
+import os
+from pathlib import Path
+from typing import Iterable
+
+# Used by Nadi (`agents/nadi_classifier.py`) to reject bad generated labels and
+# by Sabina (`agents/sabina_evaluator.py`) to score per-label accuracy.
+LABELS = ("up", "down", "neutral")
+
+# Used by Sabina's prediction validation; probabilities are rounded, so the
+# contract allows a small tolerance instead of requiring exact sums.
+PROBABILITY_TOLERANCE = 0.02
+
+# Handoff 2: Nadi writes these columns in predictions_test.csv; Sabina reads
+# them, and Nadi's LLM-code guardrail checks generated scripts against them.
+PREDICTION_COLUMNS = [
+    "article_id", "date", "ticker", "article_title", "price_t", "price_t1",
+    "pct_change", "label", "predicted_label", "confidence",
+    "prob_up", "prob_down", "prob_neutral", "split",
+]
+
+# Handoff 4: Jack writes these columns in sample_for_explanation.csv; Freddi
+# reads them before generating explanations.
+EXPLANATION_SAMPLE_COLUMNS = [
+    "article_id", "article_title", "predicted_label", "actual_label",
+    "confidence", "prob_up", "prob_down", "prob_neutral",
+]
+
+# Handoff 5: Freddi writes these columns in explanations.csv; Jack reads them
+# when building the final outputs.
+EXPLANATION_OUTPUT_COLUMNS = [
+    "article_id", "article_title", "predicted_label", "actual_label",
+    "confidence", "explanation", "manual_score",
+]
+
+# Freddi copies these input columns straight through to explanations.csv.
+EXPLANATION_PASSTHROUGH_COLUMNS = [
+    "article_id", "article_title", "predicted_label", "actual_label", "confidence",
+]
+
+# Handoff 6: Jack writes these columns in final_results.csv for the notebook and
+# any Streamlit app to read.
+FINAL_RESULTS_COLUMNS = [
+    "article_id", "date", "ticker", "article_title", "price_t", "price_t1",
+    "pct_change", "label", "predicted_label", "confidence", "explanation", "manual_score",
+]
+
+
+def read_prediction_rows(path: str | os.PathLike[str]) -> list[dict]:
+    """Read Nadi's predictions_test.csv for Sabina and enforce column order."""
+    with open(path, newline="", encoding="utf-8") as f:
+        reader = csv.DictReader(f)
+        rows = list(reader)
+    if reader.fieldnames != PREDICTION_COLUMNS:
+        raise ValueError(
+            "predictions_test.csv columns do not match data contract: "
+            f"{reader.fieldnames}"
+        )
+    if not rows:
+        raise ValueError("predictions_test.csv must contain at least one held-out row")
+    return rows
+
+
+def validate_prediction_rows(rows: Iterable[dict]) -> None:
+    """Validate Nadi prediction rows before Sabina scores them."""
+    for row in rows:
+        article_id = row.get("article_id", "")
+        label = row.get("label", "")
+        predicted = row.get("predicted_label", "")
+        if row.get("split") not in ("val", "test"):
+            raise ValueError(f"{article_id}: split must be val or test")
+        if label not in LABELS:
+            raise ValueError(f"{article_id}: invalid label {label!r}")
+        if predicted not in LABELS:
+            raise ValueError(f"{article_id}: invalid predicted_label {predicted!r}")
+
+        probs = [float(row[f"prob_{label_name}"]) for label_name in LABELS]
+        confidence = float(row["confidence"])
+        if not 0 <= confidence <= 1:
+            raise ValueError(f"{article_id}: confidence must be between 0 and 1")
+        if any(prob < 0 or prob > 1 for prob in probs):
+            raise ValueError(f"{article_id}: probabilities must be between 0 and 1")
+        if abs(sum(probs) - 1.0) > PROBABILITY_TOLERANCE:
+            raise ValueError(f"{article_id}: prob_* columns must sum to about 1")
+        if abs(confidence - max(probs)) > PROBABILITY_TOLERANCE:
+            raise ValueError(f"{article_id}: confidence must equal max prob_*")
+
+
+def _test_rows(predictions):
+    """Return test rows Jack uses for explanations and final outputs."""
+    if "split" not in predictions.columns:
+        raise ValueError("predictions file has no split column (Handoff 2)")
+    test = predictions[predictions["split"] == "test"]
+    if test.empty:
+        raise ValueError("predictions file has no split=test rows")
+    return test
+
+
+def build_explanation_sample(predictions, sample_size: int = 300):
+    """Shape Jack's sample_for_explanation.csv rows from Nadi predictions."""
+    predictions = _test_rows(predictions)
+    sample = (predictions[[
+        "article_id", "article_title", "predicted_label", "label",
+        "confidence", "prob_up", "prob_down", "prob_neutral",
+    ]].rename(columns={"label": "actual_label"}))
+    n = min(len(sample), sample_size)
+    if n < len(sample):
+        sample = sample.sample(n=n, random_state=42)
+    return sample
+
+
+def write_explanation_sample(
+    predictions_path: str | os.PathLike[str],
+    output_path: str | os.PathLike[str],
+    sample_size: int = 300,
+) -> int:
+    """Write Jack's sample_for_explanation.csv for Freddi."""
+    import pandas as pd
+
+    preds = pd.read_csv(predictions_path)
+    sample = build_explanation_sample(preds, sample_size)
+    output = Path(output_path)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    sample.to_csv(output, index=False)
+    return len(sample)
+
+
+def build_final_results(predictions, explanations):
+    """Shape Jack's final_results.csv from Nadi predictions and Freddi output."""
+    predictions = _test_rows(predictions)
+    expl = explanations[["article_id", "explanation", "manual_score"]]
+    final = predictions.merge(expl, on="article_id", how="left")[FINAL_RESULTS_COLUMNS]
+    final["explanation"] = final["explanation"].fillna("")
+    final["manual_score"] = final["manual_score"].astype("Int64")
+    return final

agents/env.py

@@ -0,0 +1,20 @@
+"""Tiny .env loader shared by the entry points (main.py, finetune_finbert.py).
+
+Populates os.environ from a local .env (KEY=VALUE lines) if it exists, so
+secrets like HF_TOKEN reach the agents — and any subprocess they spawn, which
+inherits this process's env — without the user exporting them by hand. Real
+environment variables take precedence over .env values.
+"""
+
+import os
+
+
+def load_dotenv(path=".env"):
+    if not os.path.exists(path):
+        return
+    with open(path, encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if line and not line.startswith("#") and "=" in line:
+                key, _, val = line.partition("=")
+                os.environ.setdefault(key.strip(), val.strip().strip('"').strip("'"))

agents/finetune_finbert.py

@@ -34,19 +34,13 @@
 LABELS = ["up", "down", "neutral"]
 LABEL_TO_ID = {"up": 0, "down": 1, "neutral": 2}
 
-def load_dotenv(path=".env"):
-    """Populate os.environ from a local .env (KEY=VALUE lines) if it exists —
-    same helper main.py uses, repeated here because this script runs standalone,
-    outside the pipeline entry point. Real environment variables take precedence."""
-    if not os.path.exists(path):
-        return
-    with open(path, encoding="utf-8") as f:
-        for line in f:
-            line = line.strip()
-            if line and not line.startswith("#") and "=" in line:
-                key, _, val = line.partition("=")
-                os.environ.setdefault(key.strip(), val.strip().strip('"').strip("'"))
-
+# Same .env loader main.py uses — imported here too because this script runs
+# standalone, outside the pipeline entry point. Dual import: works as a package
+# member and as a bare script (`uv run agents/finetune_finbert.py`).
+try:
+    from agents.env import load_dotenv
+except ModuleNotFoundError:
+    from env import load_dotenv
 
 load_dotenv()
 
@@ -93,19 +87,26 @@ def pick_device():
 def split_frames(df):
     """Split the data into train / validation / test DataFrames.
 
-    train/test come straight from Aurora's `split` column. We carve the
-    validation set out of the LAST 10% of training DATES (not random rows) so
-    validation looks like the real future test set and nothing leaks."""
+    All three come straight from Aurora's `split` column — she owns the one
+    date-based boundary (val = last 10% of training dates), so this script and
+    the retune loop validate on the same rows. Older files without `val` rows
+    get the same carve computed here as a fallback."""
     if "split" not in df.columns:
         raise SystemExit(
             "processed_data.csv has no 'split' column. Regenerate it with the "
             "Processing agent first: uv run agents/aurora_processing.py"
         )
 
-    train_all = df[df["split"] == "train"]
     test = df[df["split"] == "test"]
 
-    # The date that sits at the 90% mark of the training dates.
+    if (df["split"] == "val").any():
+        train = df[df["split"] == "train"]
+        val = df[df["split"] == "val"]
+        return train, val, test
+
+    # Fallback for pre-val files: carve the LAST 10% of training DATES (not
+    # random rows) so validation looks like the real future test set.
+    train_all = df[df["split"] == "train"]
     val_cutoff = pd.to_datetime(train_all["date"]).quantile(0.9)
     is_val = pd.to_datetime(train_all["date"]) > val_cutoff
 

agents/freddi_explanation.py

@@ -42,22 +42,18 @@
 # sys.path) — same dual-import trick the other agents use.
 try:
     from agents.base import Agent
+    from agents.contracts import (
+        EXPLANATION_OUTPUT_COLUMNS as OUTPUT_COLUMNS,
+        EXPLANATION_PASSTHROUGH_COLUMNS as PASSTHROUGH_COLUMNS,
+        EXPLANATION_SAMPLE_COLUMNS as INPUT_COLUMNS,
+    )
 except ModuleNotFoundError:
     from base import Agent
-
-# Columns read in (Handoff 4) and the exact columns written out (Handoff 5).
-INPUT_COLUMNS = [
-    "article_id", "article_title", "predicted_label", "actual_label",
-    "confidence", "prob_up", "prob_down", "prob_neutral",
-]
-OUTPUT_COLUMNS = [
-    "article_id", "article_title", "predicted_label", "actual_label",
-    "confidence", "explanation", "manual_score",
-]
-# Passed through from the input untouched (Golden rule 1).
-PASSTHROUGH_COLUMNS = [
-    "article_id", "article_title", "predicted_label", "actual_label", "confidence",
-]
+    from contracts import (
+        EXPLANATION_OUTPUT_COLUMNS as OUTPUT_COLUMNS,
+        EXPLANATION_PASSTHROUGH_COLUMNS as PASSTHROUGH_COLUMNS,
+        EXPLANATION_SAMPLE_COLUMNS as INPUT_COLUMNS,
+    )
 
 DEFAULT_INPUT = "mock_data/sample_for_explanation.csv"
 DEFAULT_OUTPUT = "explanations.csv"