"""Unit tests for `lexsi_ds.agent.tools.predict`. Most tests run after a successful `train_tabular_model` (so model_id / predict tag context is in cache). A few isolate the negative paths. """ from __future__ import annotations import pandas as pd import pytest from lexsi_ds.agent.tools.predict import ( TOOL, PredictArgs, _align_inference_columns, _first_present, _mint_predict_tag, _normalize_predictions, _take_top_k, ) from lexsi_ds.agent.tools.train_tabular_model import ( TOOL as TRAIN_TOOL, TrainTabularModelArgs, ) def _train_first(ctx, df, target="status", entity="loan_id"): ctx.cache["sql_result:ctx_df"] = df ctx.cache["last_sql_result"] = df out = TRAIN_TOOL.run( TrainTabularModelArgs( df_label="ctx_df", target_column=target, entity_column=entity, ), ctx, ) assert out.ok, out.summary return out # ----- happy path ----- def test_happy_path(ctx_with_fake, loan_context_df, loan_predict_df): _train_first(ctx_with_fake, loan_context_df) ctx_with_fake.cache["sql_result:predict_df"] = loan_predict_df out = TOOL.run( PredictArgs(df_label="predict_df", top_k=5), ctx_with_fake, ) assert out.ok, out.summary assert "Prediction" in out.payload["df"].columns assert out.payload["n_rows"] == len(loan_predict_df) assert out.payload["pred_col"] == "Prediction" # Cache contract assert "last_predictions" in ctx_with_fake.cache assert ctx_with_fake.cache["last_predict_tag"].startswith("agentpredict") assert ctx_with_fake.cache["last_predict_df_label"] == "predict_df" cols = ctx_with_fake.cache["last_prediction_columns"] assert cols["pred"] == "Prediction" assert cols["entity"] == "loan_id" def test_explicit_model_id_passthrough(ctx_with_fake, loan_context_df, loan_predict_df): train_out = _train_first(ctx_with_fake, loan_context_df) model_id = train_out.payload["model_id"] ctx_with_fake.cache["sql_result:predict_df"] = loan_predict_df # Wipe the cache last_model_id to force the explicit-arg path ctx_with_fake.cache["last_model_id"] = None out = TOOL.run( PredictArgs(df_label="predict_df", model_id=model_id), ctx_with_fake, ) assert out.ok, out.summary assert out.payload["model_id"] == model_id def test_top_k_observation_is_bounded(ctx_with_fake, loan_context_df, loan_predict_df): _train_first(ctx_with_fake, loan_context_df) ctx_with_fake.cache["sql_result:predict_df"] = loan_predict_df out = TOOL.run( PredictArgs(df_label="predict_df", top_k=3), ctx_with_fake, ) assert out.ok assert len(out.payload["top_k_df"]) == 3 def test_strips_null_target_column(ctx_with_fake, loan_context_df, loan_predict_df): # loan_predict_df already has all-null status — confirm it doesn't blow up _train_first(ctx_with_fake, loan_context_df) ctx_with_fake.cache["sql_result:predict_df"] = loan_predict_df out = TOOL.run( PredictArgs(df_label="predict_df"), ctx_with_fake, ) assert out.ok, out.summary # ----- error branches ----- def test_missing_df(ctx_with_fake, loan_context_df): _train_first(ctx_with_fake, loan_context_df) out = TOOL.run(PredictArgs(df_label="not_present"), ctx_with_fake) assert not out.ok assert out.error == "missing_df" def test_empty_predict_df(ctx_with_fake, loan_context_df): _train_first(ctx_with_fake, loan_context_df) ctx_with_fake.cache["sql_result:predict_df"] = pd.DataFrame( columns=["loan_id", "amount"] ) out = TOOL.run(PredictArgs(df_label="predict_df"), ctx_with_fake) assert not out.ok assert out.error == "empty_predict_df" def test_no_model_id_no_cache(ctx_with_fake, loan_predict_df): ctx_with_fake.cache["sql_result:predict_df"] = loan_predict_df out = TOOL.run(PredictArgs(df_label="predict_df"), ctx_with_fake) assert not out.ok assert out.error == "no_model_id" def test_no_tab_project(ctx_offline, loan_predict_df): ctx_offline.cache["sql_result:predict_df"] = loan_predict_df ctx_offline.cache["last_model_id"] = "fake_model" out = TOOL.run(PredictArgs(df_label="predict_df"), ctx_offline) assert not out.ok assert out.error == "no_tab_project" # ----- helpers ----- def test_first_present(): cols = pd.Index(["a", "Prediction", "x"]) assert _first_present(cols, ("pred", "Prediction", "y")) == "Prediction" assert _first_present(cols, ("z",)) is None def test_take_top_k_with_prob(): df = pd.DataFrame({"id": [1, 2, 3, 4], "prob": [0.1, 0.9, 0.5, 0.7]}) top = _take_top_k(df, prob_col="prob", pred_col="id", k=2, task_type="classification") assert top["id"].tolist() == [2, 4] def test_take_top_k_regression_abs(): df = pd.DataFrame({"id": [1, 2, 3, 4], "pred": [-3, 1, 0.5, -10]}) top = _take_top_k(df, prob_col=None, pred_col="pred", k=2, task_type="regression") assert top["id"].tolist() == [4, 1] def test_normalize_predictions_attaches_entity(loan_predict_df): raw = pd.DataFrame({ "Prediction": ["repaid"] * len(loan_predict_df), "Probability": [0.5] * len(loan_predict_df), }) out, pred_c, prob_c = _normalize_predictions( preds_raw=raw, predict_df=loan_predict_df, entity_col="loan_id", target_col="status", ) assert "loan_id" in out.columns assert pred_c == "Prediction" assert prob_c == "Probability" def test_mint_predict_tag_is_safe(): """Lexsi rejects ANY non-alphanumeric in tag names (including `_`). Strip everything except [a-zA-Z0-9]; uniqueness comes from run_id.""" t = _mint_predict_tag("r1", "weird/label!") assert t.isalnum(), f"non-alphanumeric chars in tag: {t!r}" assert t.startswith("agentpredict") assert "r1" in t # run_id preserved # ---------- inference-frame schema alignment (Lexsi #05-018) ---------- def test_align_inference_columns_adds_missing_and_drops_extra(): """Missing training features are filled with '' and extras dropped, in training order, with the target excluded — preventing SDK #05-018.""" train_cols = ["loan_id", "loan_duration", "n_trans_all", "n_orders", "defaulted"] pdf = pd.DataFrame({ "loan_id": [1, 2], "loan_duration": [12, 24], "district": ["A", "B"], # extra column the model never saw }) aligned, added, dropped = _align_inference_columns(pdf, train_cols, target_col="defaulted") assert list(aligned.columns) == ["loan_id", "loan_duration", "n_trans_all", "n_orders"] assert added == ["n_trans_all", "n_orders"] assert dropped == ["district"] assert (aligned["n_trans_all"] == "").all() assert (aligned["n_orders"] == "").all() def test_align_inference_columns_noop_without_schema(): """No cached training schema → frame returned untouched (e.g. predicting against a model not trained in this session).""" pdf = pd.DataFrame({"a": [1], "b": [2]}) aligned, added, dropped = _align_inference_columns(pdf, None, target_col="t") assert aligned is pdf assert added == [] and dropped == []