|
| 1 | +import json |
| 2 | +import os |
| 3 | +from pathlib import Path |
| 4 | + |
| 5 | +import pandas as pd |
| 6 | +import requests |
| 7 | +from dotenv import load_dotenv |
| 8 | +from langchain.agents import create_agent |
| 9 | +from langchain.tools import tool |
| 10 | +from langchain_litellm import ChatLiteLLM |
| 11 | + |
| 12 | +from payload_data import payload, GROUND_TRUTH |
| 13 | + |
| 14 | +load_dotenv(dotenv_path=Path(__file__).parent / ".env") |
| 15 | + |
| 16 | +TABPFN_DEPLOYMENT_ID = os.environ["TABPFN_DEPLOYMENT_ID"] |
| 17 | + |
| 18 | +model = ChatLiteLLM(model=os.environ.get("LLM_MODEL", "sap/amazon--nova-lite"), temperature=0) |
| 19 | + |
| 20 | + |
| 21 | +def _get_auth_token() -> str: |
| 22 | + response = requests.post( |
| 23 | + os.environ["AICORE_AUTH_URL"], |
| 24 | + data={ |
| 25 | + "grant_type": "client_credentials", |
| 26 | + "client_id": os.environ["AICORE_CLIENT_ID"], |
| 27 | + "client_secret": os.environ["AICORE_CLIENT_SECRET"], |
| 28 | + }, |
| 29 | + timeout=30, |
| 30 | + ) |
| 31 | + response.raise_for_status() |
| 32 | + return response.json()["access_token"] |
| 33 | + |
| 34 | + |
| 35 | +FEATURE_COLS = [ |
| 36 | + "CATEGORY", "DISTANCE_KM", "WEIGHT_KG", "ORDER_VALUE", "PROMISED_DAYS", |
| 37 | + "WAREHOUSE_HRS", "SHIP_METHOD", "CARRIER", "PKG_SIZE", "PRIORITY", "WEATHER", "SHIP_COST", |
| 38 | +] |
| 39 | + |
| 40 | + |
| 41 | +def call_tabpfn(payload: dict) -> str: |
| 42 | + try: |
| 43 | + token = _get_auth_token() |
| 44 | + url = f"{os.environ['AICORE_BASE_URL'].rstrip('/')}/v2/inference/deployments/{TABPFN_DEPLOYMENT_ID}/predict" |
| 45 | + headers = { |
| 46 | + "Authorization": f"Bearer {token}", |
| 47 | + "AI-Resource-Group": os.environ.get("AICORE_RESOURCE_GROUP", "default"), |
| 48 | + "Content-Type": "application/json", |
| 49 | + } |
| 50 | + |
| 51 | + rows = payload["rows"] |
| 52 | + results: dict = {"predictions": {}} |
| 53 | + |
| 54 | + for target in payload["prediction_config"]["target_columns"]: |
| 55 | + col_name = target["name"] |
| 56 | + placeholder = target["prediction_placeholder"] |
| 57 | + task_type = target["task_type"] |
| 58 | + |
| 59 | + # .get() handles classification-only test rows that have no DELAY_DAYS key |
| 60 | + train_rows = [r for r in rows if r.get(col_name) is not None and r.get(col_name) != placeholder] |
| 61 | + test_rows = [r for r in rows if r.get(col_name) == placeholder] |
| 62 | + |
| 63 | + if not test_rows: |
| 64 | + results["predictions"][col_name] = [] |
| 65 | + continue |
| 66 | + |
| 67 | + train_df = pd.DataFrame(train_rows) |
| 68 | + test_df = pd.DataFrame(test_rows) |
| 69 | + |
| 70 | + x_train = {col: train_df[col].tolist() for col in FEATURE_COLS} |
| 71 | + x_test = {col: test_df[col].tolist() for col in FEATURE_COLS} |
| 72 | + |
| 73 | + if task_type == "regression": |
| 74 | + request_payload = { |
| 75 | + "task_config": { |
| 76 | + "task": "regression", |
| 77 | + "predict_params": {"output_type": "mean"}, |
| 78 | + }, |
| 79 | + "x_train": x_train, |
| 80 | + "y_train": train_df[col_name].astype(float).tolist(), |
| 81 | + "x_test": x_test, |
| 82 | + } |
| 83 | + else: |
| 84 | + request_payload = { |
| 85 | + "task_config": { |
| 86 | + "task": "classification", |
| 87 | + "tabpfn_config": {"n_estimators": 8, "inference_precision": "auto"}, |
| 88 | + "predict_params": {"output_type": "preds"}, |
| 89 | + }, |
| 90 | + "x_train": x_train, |
| 91 | + "y_train": {col_name: train_df[col_name].tolist()}, |
| 92 | + "x_test": x_test, |
| 93 | + } |
| 94 | + |
| 95 | + print(f"\n--- TabPFN request [{col_name}] ---") |
| 96 | + print(f"URL: {url}") |
| 97 | + response = requests.post(url, headers=headers, json=request_payload, timeout=300) |
| 98 | + print(f"HTTP {response.status_code}") |
| 99 | + response.raise_for_status() |
| 100 | + |
| 101 | + predictions = response.json()["prediction"] |
| 102 | + results["predictions"][col_name] = [ |
| 103 | + {"ORD_ID": row["ORD_ID"], "SHIP_METHOD": row["SHIP_METHOD"], col_name: p} |
| 104 | + for row, p in zip(test_rows, predictions) |
| 105 | + ] |
| 106 | + |
| 107 | + return json.dumps(results, indent=2) |
| 108 | + |
| 109 | + except requests.RequestException as e: |
| 110 | + return f"Error calling TabPFN: {str(e)}" |
| 111 | + |
| 112 | + |
| 113 | +@tool |
| 114 | +def run_tabpfn_predictions() -> str: |
| 115 | + """Run TabPFN-3.5 Plus predictions for shipment delay days (DELAY_DAYS) and on-time status (LATE). |
| 116 | +
|
| 117 | + Returns a JSON string with predictions for all test rows. |
| 118 | + """ |
| 119 | + return call_tabpfn(payload) |
| 120 | + |
| 121 | + |
| 122 | +system_prompt = """You are a senior e-commerce logistics analyst. |
| 123 | +You have received TabPFN-3.5 Plus predictions for two missing fields in a shipment dataset: |
| 124 | +- DELAY_DAYS (integer): the predicted number of days an order arrives after its promised delivery date. |
| 125 | +- LATE (Yes/No): whether each order is predicted to arrive after its promised delivery date. |
| 126 | +
|
| 127 | +Your role is to summarise the predictions in a clear operations report: |
| 128 | +- List the predicted delivery delay (in days) for each order, grouped by shipping method. |
| 129 | +- List the predicted on-time / late status for each order. |
| 130 | +- Highlight any patterns (e.g. which shipping methods, carriers, or weather conditions correlate with longer delays). |
| 131 | +
|
| 132 | +You rely exclusively on the model predictions provided — never estimate or adjust values yourself.""" |
| 133 | + |
| 134 | + |
| 135 | +forecast_agent = create_agent( |
| 136 | + model=model, |
| 137 | + tools=[run_tabpfn_predictions], |
| 138 | + system_prompt=system_prompt, |
| 139 | +) |
| 140 | + |
| 141 | + |
| 142 | +def evaluate_predictions(predictions_json: str) -> None: |
| 143 | + data = json.loads(predictions_json) |
| 144 | + preds = data["predictions"] |
| 145 | + |
| 146 | + if "DELAY_DAYS" in preds: |
| 147 | + errors = [] |
| 148 | + for item in preds["DELAY_DAYS"]: |
| 149 | + truth = GROUND_TRUTH.get(item["ORD_ID"], {}).get("DELAY_DAYS") |
| 150 | + if truth is not None: |
| 151 | + errors.append(abs(float(item["DELAY_DAYS"]) - float(truth))) |
| 152 | + if errors: |
| 153 | + print(f"\nMAE (DELAY_DAYS): {sum(errors) / len(errors):.2f} days") |
| 154 | + |
| 155 | + if "LATE" in preds: |
| 156 | + correct = 0 |
| 157 | + total = 0 |
| 158 | + for item in preds["LATE"]: |
| 159 | + truth = GROUND_TRUTH.get(item["ORD_ID"], {}).get("LATE") |
| 160 | + if truth is not None: |
| 161 | + total += 1 |
| 162 | + if str(item["LATE"]).strip().lower() == str(truth).strip().lower(): |
| 163 | + correct += 1 |
| 164 | + if total: |
| 165 | + print(f"Accuracy (LATE): {correct}/{total} = {correct / total:.0%}") |
| 166 | + |
| 167 | + |
| 168 | +def main(): |
| 169 | + result = forecast_agent.invoke( |
| 170 | + {"messages": [{"role": "user", "content": "Run the TabPFN predictions and write a concise logistics operations summary."}]} |
| 171 | + ) |
| 172 | + |
| 173 | + print("\n" + "=" * 50) |
| 174 | + print("E-commerce Logistics Report:") |
| 175 | + print("=" * 50) |
| 176 | + print(result["messages"][-1].content) |
| 177 | + |
| 178 | + for msg in result["messages"]: |
| 179 | + if hasattr(msg, "content") and isinstance(msg.content, str): |
| 180 | + try: |
| 181 | + data = json.loads(msg.content) |
| 182 | + if "predictions" in data: |
| 183 | + evaluate_predictions(msg.content) |
| 184 | + break |
| 185 | + except (json.JSONDecodeError, KeyError): |
| 186 | + pass |
| 187 | + |
| 188 | + |
| 189 | +if __name__ == "__main__": |
| 190 | + main() |
0 commit comments