drift-ml 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

drift/cli/client.py

@@ -123,3 +123,17 @@ class BackendClient:
                 detail = r.text or str(r.status_code)
             raise BackendError(f"Train failed: {detail}", status_code=r.status_code, body=r.text)
         return r.json()
+
+    def download_notebook(self, run_id: str) -> Optional[bytes]:
+        """GET /download/{run_id}/notebook — download notebook bytes, or None if failed."""
+        r = requests.get(self._url(f"/download/{run_id}/notebook"), timeout=60)
+        if r.status_code != 200:
+            return None
+        return r.content
+
+    def download_model(self, run_id: str) -> Optional[bytes]:
+        """GET /download/{run_id}/model — download model pickle bytes, or None if refused/failed."""
+        r = requests.get(self._url(f"/download/{run_id}/model"), timeout=60)
+        if r.status_code != 200:
+            return None
+        return r.content

drift/cli/repl.py

@@ -3,10 +3,12 @@ Chat-based CLI loop for drift.
 Natural language input; maintains session state; reuses backend planner + executor.
 """
 
+import os
 import re
 import sys
 import threading
 import time
+from pathlib import Path
 from typing import Any, Dict, Optional
 
 from drift.cli.client import BackendClient, BackendError, DEFAULT_BASE_URL
@@ -194,8 +196,9 @@ def _run_training_and_show(client: BackendClient, session: SessionState) -> None
         print(f"Training error: {train_error}", file=sys.stderr)
         return
     if train_result:
+        run_id = train_result.get("run_id", "")
         session.update_after_train(
-            run_id=train_result.get("run_id", ""),
+            run_id=run_id,
             metrics=train_result.get("metrics") or {},
             agent_message=train_result.get("agent_message"),
         )
@@ -207,6 +210,37 @@ def _run_training_and_show(client: BackendClient, session: SessionState) -> None
             primary = metrics.get("primary_metric_value") or metrics.get("accuracy") or metrics.get("r2")
             if primary is not None:
                 print(f"  Primary metric: {primary}")
+        # Auto-save notebook and model to current directory for terminal users
+        if run_id and not train_result.get("refused"):
+            _save_artifacts(client, run_id)
+
+
+def _save_artifacts(client: BackendClient, run_id: str) -> None:
+    """Download notebook and model to current directory. Print where saved."""
+    cwd = Path(os.getcwd())
+    prefix = f"model_{run_id[:8]}"
+    saved = []
+    try:
+        nb = client.download_notebook(run_id)
+        if nb:
+            path = cwd / f"{prefix}_training_notebook.ipynb"
+            path.write_bytes(nb)
+            saved.append(str(path))
+    except Exception:
+        pass
+    try:
+        mdl = client.download_model(run_id)
+        if mdl:
+            path = cwd / f"{prefix}.pkl"
+            path.write_bytes(mdl)
+            saved.append(str(path))
+    except Exception:
+        pass
+    if saved:
+        print()
+        print("  Saved to current directory:")
+        for p in saved:
+            print(f"    {p}")
 
 
 def _print_message(content: str) -> None:

{drift_ml-0.1.1.dist-info → drift_ml-0.1.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: drift-ml
-Version: 0.1.1
+Version: 0.1.3
 Summary: Terminal-first AutoML CLI - chat-based ML engineer
 Requires-Python: >=3.10
 Requires-Dist: requests>=2.28.0

{drift_ml-0.1.1.dist-info → drift_ml-0.1.3.dist-info}/RECORD

@@ -1,15 +1,15 @@
 drift/__init__.py,sha256=X0NUP5ZAZSz-rBFfjvmmS4IYWP2_CZtu187mqwpWwqk,127
 drift/__main__.py,sha256=MMUjNUbctbLHVe37ZCOW-66h8OZ8WIYiISxHJHRNxes,202
 drift/cli/__init__.py,sha256=OQt7M06e98e_L_60Qz-HsmSqGXKWdH53SvgETJ_BZZ0,172
-drift/cli/client.py,sha256=AiV5hNCWgQbH9biznPUD7jctsSpDaXdy1WcgAIoBqtM,4880
-drift/cli/repl.py,sha256=ifZH2xoWFwCfdsGSh7HEkdO378t1c4_h07iNZGiUrMs,7220
+drift/cli/client.py,sha256=hGxz-J8fC_gGV9BZnpCBwkQZg4Q0_s9-Agk2PuXQ5MQ,5527
+drift/cli/repl.py,sha256=hRnx4nfP_StnXDbdRZXnnjA8aNvrdQSV-uTNsuaaO60,8262
 drift/cli/session.py,sha256=fMBI_pgiOmm-hcf7GDcjDeMIKESy6mGeEWhWyzdBirM,2860
 drift/llm_adapters/__init__.py,sha256=y1UhZWlC8Ik_OKfLcOp0JZP-FKR3MBBCemWwsL6TnkY,296
 drift/llm_adapters/base.py,sha256=KlZUPYpvCI8pafklBWel0GHHLNtVKVwSHA92MZt3VsI,1331
 drift/llm_adapters/gemini_cli.py,sha256=Z61wY3yFiZqPrQrpJAQrBtMbBkr2qLWBkni-A4p9lZo,2163
 drift/llm_adapters/local_llm.py,sha256=Z6j6z1CXk2LMeQ5ZnY4o38PiYkHYmcIkgJGkHdU50M8,2279
-drift_ml-0.1.1.dist-info/METADATA,sha256=EGmvqr1xDEnVM42RZFtYS4htDu7KfDimXC2aD9qPu6o,168
-drift_ml-0.1.1.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-drift_ml-0.1.1.dist-info/entry_points.txt,sha256=aCY7U9M8nhYj_tIfTXJmYkVmXY3ZoxF0tebDZzYswv8,46
-drift_ml-0.1.1.dist-info/top_level.txt,sha256=3u2KGqsciGZQ2uCoBivm55t3e8er8S4xnqkgdQ_8oeM,6
-drift_ml-0.1.1.dist-info/RECORD,,
+drift_ml-0.1.3.dist-info/METADATA,sha256=JYHG8KKtZxOeOrkaCUqZ2WnRDtxYnRqoGtDE9REyKmI,168
+drift_ml-0.1.3.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+drift_ml-0.1.3.dist-info/entry_points.txt,sha256=aCY7U9M8nhYj_tIfTXJmYkVmXY3ZoxF0tebDZzYswv8,46
+drift_ml-0.1.3.dist-info/top_level.txt,sha256=3u2KGqsciGZQ2uCoBivm55t3e8er8S4xnqkgdQ_8oeM,6
+drift_ml-0.1.3.dist-info/RECORD,,