Spaces:
Sleeping
Sleeping
| """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 == [] | |