A local-first private AI assistant for everyday use. Runs on-device models with encrypted P2P sync, and supports sharing chats publicly on ATProto.
Merge pull request #16 from tilesprivacy/feat/model-auto-download
feat: auto downloading model given in modelfile
authored by
Anandu Pavanan
and committed by
GitHub
2e913ad3
fab93cea
+354
-8
+18
-5
server/api.py
+18
-5
server/api.py
···
35
35
from .cache_utils import (
36
36
get_model_path
37
37
)
38
+
from .hf_downloader import pull_model
39
+
38
40
from .mlx_runner import MLXRunner
39
41
40
42
from server.mem_agent.utils import extract_python_code, extract_reply, extract_thoughts, create_memory_if_not_exists, format_results
···
107
109
class StartRequest(BaseModel):
108
110
model: str
109
111
memory_path: str
112
+
113
+
class downloadRequest(BaseModel):
114
+
model: str
110
115
111
116
class Agent:
112
117
def __init__(
···
198
203
async def ping():
199
204
return {"message": "Badda-Bing Badda-Bang"}
200
205
206
+
@app.post("/download")
207
+
async def download(request:downloadRequest):
208
+
""" Download the model """
209
+
try:
210
+
if pull_model(request.model):
211
+
return {"message": "Model downloaded"}
212
+
else:
213
+
raise HTTPException(status_code=400, detail="Downloading model failed")
214
+
except Exception as e:
215
+
raise HTTPException(status_code=500, detail=str(e))
216
+
201
217
@app.post("/start")
202
218
async def start_model(request: StartRequest):
203
219
"""Load the model and start the agent"""
···
206
222
_messages = [ChatMessage(role="system", content=SYSTEM_PROMPT)]
207
223
_memory_path = request.memory_path
208
224
209
-
try:
210
-
_runner = get_or_load_model(request.model)
211
-
return {"message": "Model loaded"}
212
-
except Exception as e:
213
-
raise HTTPException(status_code=500, detail=str(e))
225
+
_runner = get_or_load_model(request.model)
226
+
return {"message": "Model loaded"}
214
227
215
228
@app.post("/v1/chat/completions")
216
229
async def create_chat_completion(request: ChatCompletionRequest):
+145
server/hf_downloader.py
+145
server/hf_downloader.py
···
1
+
import json
2
+
import os
3
+
import subprocess
4
+
import sys
5
+
import tempfile
6
+
7
+
try:
8
+
from .cache_utils import (
9
+
MODEL_CACHE,
10
+
hf_to_cache_dir,
11
+
is_model_healthy,
12
+
parse_model_spec,
13
+
)
14
+
except ImportError:
15
+
from pathlib import Path
16
+
def parse_model_spec(x): return (x, None)
17
+
def hf_to_cache_dir(x): return x
18
+
if "HF_HOME" in os.environ:
19
+
MODEL_CACHE = Path(os.environ["HF_HOME"]) / "hub"
20
+
else:
21
+
MODEL_CACHE = Path(os.path.expanduser("~/.cache/huggingface/hub"))
22
+
def is_model_healthy(x): return False
23
+
24
+
def describe_http_exception(exc):
25
+
if hasattr(exc, "response") and exc.response is not None:
26
+
status = getattr(exc.response, "status_code", None)
27
+
url = getattr(exc.response, "url", None)
28
+
if status == 401:
29
+
return f"[ERROR] Unauthorized (401): Check your HuggingFace token or login.\nURL: {url}"
30
+
elif status == 403:
31
+
return f"[ERROR] Forbidden (403): Access denied.\nURL: {url}"
32
+
elif status == 404:
33
+
return f"[ERROR] Not Found (404): Resource does not exist.\nURL: {url}"
34
+
elif status >= 500:
35
+
return f"[ERROR] Server Error ({status}): Problem on HuggingFace's side.\nURL: {url}\nTry again later."
36
+
else:
37
+
return f"[ERROR] HTTP Error {status}: {exc}\nURL: {url}"
38
+
return f"[ERROR] HTTP Error: {exc}"
39
+
40
+
def configure_download_environment():
41
+
os.environ['HF_HUB_DOWNLOAD_THREADS'] = '1'
42
+
os.environ['HF_HUB_DOWNLOAD_CHUNK_SIZE'] = '524288' # 512KB chunks for household-friendly downloads
43
+
os.environ['HF_HUB_ENABLE_HF_TRANSFER'] = 'false'
44
+
45
+
def pull_model(model_spec):
46
+
original_spec = model_spec
47
+
model_name, commit_hash = parse_model_spec(model_spec)
48
+
49
+
# Validate HuggingFace Hub repository name length limit (Issue #6)
50
+
if len(model_name) > 96:
51
+
print(f"[ERROR] Repository name exceeds HuggingFace Hub limit: {len(model_name)}/96 characters")
52
+
print("Repository names longer than 96 characters cannot exist on HuggingFace Hub.")
53
+
print(f"Invalid name: '{model_name}'")
54
+
return False
55
+
56
+
if "/" not in original_spec.split("@")[0] and "/" in model_name:
57
+
confirm = input(f"Download '{model_name}'? [Y/n] ")
58
+
if confirm.lower() == "n":
59
+
print("Download cancelled.")
60
+
return False
61
+
62
+
base_cache_dir = MODEL_CACHE / hf_to_cache_dir(model_name)
63
+
if commit_hash:
64
+
hash_dir = base_cache_dir / "snapshots" / commit_hash
65
+
if hash_dir.exists() and is_model_healthy(f"{model_name}@{commit_hash}"):
66
+
print("Model already exists")
67
+
return True
68
+
else:
69
+
if base_cache_dir.exists() and is_model_healthy(model_name):
70
+
print("Model already exists")
71
+
return True
72
+
73
+
print(f"Downloading {model_name}...")
74
+
75
+
# Build kwargs dict for the worker
76
+
kwargs_dict = {
77
+
"repo_id": model_name,
78
+
"local_dir_use_symlinks": False,
79
+
"max_workers": 1
80
+
}
81
+
if commit_hash:
82
+
kwargs_dict["revision"] = commit_hash
83
+
# if "mlx-community" in model_name:
84
+
kwargs_dict["allow_patterns"] = [
85
+
"*.json", "*.txt", "*.safetensors", "*.md", "*.gitattributes", "LICENSE"
86
+
]
87
+
# if "mlx-community" not in model_name:
88
+
# confirm = input(f"[WARNING] {model_name} is not an MLX model (may be >1GB). Continue? [y/N] ")
89
+
# if confirm.lower() != "y":
90
+
# print("Download cancelled.")
91
+
# return
92
+
93
+
kwargs_str = json.dumps(kwargs_dict, indent=2)
94
+
with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
95
+
f.write(kwargs_str)
96
+
kwargs_file = f.name
97
+
98
+
# Call the worker as subprocess with nice priority
99
+
worker_path = os.path.join(os.path.dirname(__file__), "throttled_download_worker.py")
100
+
try:
101
+
result = subprocess.run(
102
+
['nice', '-n', '19', sys.executable, worker_path, kwargs_file],
103
+
check=False
104
+
)
105
+
if result.returncode == 0:
106
+
print("Download completed successfully.")
107
+
return True
108
+
elif result.returncode in (10, 11, 12, 13, 14, 15):
109
+
# Already handled in worker, do NOT retry fallback
110
+
print("[WARNING] Fatal error encountered in throttled download, not attempting fallback.")
111
+
return False
112
+
else:
113
+
print("[WARNING] Throttled download failed or was interrupted.")
114
+
print("Attempting fallback download with standard throttling...")
115
+
try:
116
+
import requests
117
+
from huggingface_hub import snapshot_download
118
+
configure_download_environment()
119
+
snapshot_download(**kwargs_dict)
120
+
print("Download completed successfully.")
121
+
return True
122
+
except requests.exceptions.HTTPError as e:
123
+
print(describe_http_exception(e))
124
+
return False
125
+
except requests.exceptions.ConnectionError:
126
+
print("[ERROR] Network connection error. Please check your internet connection and try again.")
127
+
return False
128
+
except requests.exceptions.Timeout:
129
+
print("[ERROR] Download timed out. Please try again.")
130
+
return False
131
+
except KeyboardInterrupt:
132
+
print("\n[WARNING] Download cancelled by user.")
133
+
return False
134
+
except Exception as e:
135
+
print(f"[ERROR] Unexpected error during fallback download: {type(e).__name__}: {e}")
136
+
return False
137
+
except KeyboardInterrupt:
138
+
print("\n[WARNING] Download cancelled by user.")
139
+
return False
140
+
except ImportError:
141
+
print("huggingface-hub is not installed. Please install it with: pip install huggingface-hub")
142
+
return False
143
+
except Exception as e:
144
+
print(f"[ERROR] Unexpected error: {type(e).__name__}: {e}")
145
+
return False
+2
-1
server/pyproject.toml
+2
-1
server/pyproject.toml
+163
server/throttled_download_worker.py
+163
server/throttled_download_worker.py
···
1
+
import json
2
+
import os
3
+
import signal
4
+
import sys
5
+
import time
6
+
from typing import Any
7
+
8
+
# Global tracking for accurate download rate
9
+
_download_stats = {
10
+
'bytes_downloaded': 0,
11
+
'start_time': None,
12
+
'last_update': None,
13
+
'actual_download_time': 0.0 # Time spent actually downloading (without delays)
14
+
}
15
+
16
+
17
+
def signal_handler(signum: int, frame: Any) -> None:
18
+
print("\n[WARNING] Download cancelled by user.")
19
+
sys.exit(0)
20
+
21
+
signal.signal(signal.SIGINT, signal_handler)
22
+
signal.signal(signal.SIGTERM, signal_handler)
23
+
24
+
os.environ["HF_HUB_DOWNLOAD_THREADS"] = "1"
25
+
os.environ["HF_HUB_DOWNLOAD_CHUNK_SIZE"] = "524288" # 512KB chunks (half size)
26
+
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "false"
27
+
28
+
try:
29
+
import requests
30
+
from huggingface_hub import snapshot_download
31
+
except ImportError:
32
+
print("[ERROR] huggingface_hub or requests not installed in worker environment!")
33
+
sys.exit(2)
34
+
35
+
# Throttle all HTTP(S) requests with adaptive delays
36
+
original_get = requests.get
37
+
original_post = requests.post
38
+
39
+
def get_adaptive_delay(url: str, response: Any) -> float:
40
+
"""Calculate delay based on file type and size"""
41
+
if not url:
42
+
return 1.0
43
+
44
+
# Check if this is a large model file download
45
+
if any(ext in url.lower() for ext in ['.safetensors', '.bin', '.pth']):
46
+
# For large model files, use more aggressive throttling
47
+
content_length = response.headers.get('content-length')
48
+
if content_length:
49
+
size_mb = int(content_length) / (1024 * 1024)
50
+
if size_mb > 100: # Files larger than 100MB
51
+
return 3.0 # 3 second delay between chunks
52
+
elif size_mb > 10: # Files larger than 10MB
53
+
return 2.0 # 2 second delay
54
+
return 2.0 # Default for model files
55
+
56
+
# Regular files (config.json, tokenizer files, etc.)
57
+
return 0.5
58
+
59
+
def throttled_get(*args: Any, **kwargs: Any) -> Any:
60
+
download_start = time.time()
61
+
response = original_get(*args, **kwargs)
62
+
download_end = time.time()
63
+
64
+
# Track actual download time (without delays)
65
+
actual_download_time = download_end - download_start
66
+
_download_stats['actual_download_time'] += actual_download_time
67
+
68
+
# Track bytes if we can determine them
69
+
url = args[0] if args else kwargs.get('url', '')
70
+
if hasattr(response, 'headers') and 'content-length' in response.headers:
71
+
content_length = int(response.headers['content-length'])
72
+
_download_stats['bytes_downloaded'] += content_length
73
+
74
+
# Initialize timing if first download
75
+
if _download_stats['start_time'] is None:
76
+
_download_stats['start_time'] = download_start
77
+
78
+
# Print accurate rate every ~5MB or every 10 seconds
79
+
now = time.time()
80
+
if (_download_stats['last_update'] is None or
81
+
now - _download_stats['last_update'] > 10 or
82
+
_download_stats['bytes_downloaded'] % (5 * 1024 * 1024) < content_length):
83
+
84
+
if _download_stats['actual_download_time'] > 0:
85
+
real_rate_mbps = (_download_stats['bytes_downloaded'] / _download_stats['actual_download_time']) / (1024 * 1024)
86
+
total_mb = _download_stats['bytes_downloaded'] / (1024 * 1024)
87
+
print(f"[THROTTLE] Downloaded {total_mb:.1f}MB at real rate: {real_rate_mbps:.1f}MB/s (excluding delays)")
88
+
_download_stats['last_update'] = now
89
+
90
+
delay = get_adaptive_delay(url, response)
91
+
time.sleep(delay)
92
+
return response
93
+
94
+
def throttled_post(*args: Any, **kwargs: Any) -> Any:
95
+
response = original_post(*args, **kwargs)
96
+
time.sleep(0.5)
97
+
return response
98
+
99
+
requests.get = throttled_get
100
+
requests.post = throttled_post
101
+
102
+
def main() -> None:
103
+
if len(sys.argv) != 2:
104
+
print("Usage: python throttled_download_worker.py <kwargs_file.json>")
105
+
sys.exit(1)
106
+
107
+
kwargs_file = sys.argv[1]
108
+
try:
109
+
with open(kwargs_file) as f:
110
+
kwargs_dict = json.load(f)
111
+
except Exception as e:
112
+
print(f"[ERROR] Could not read worker kwargs: {e}")
113
+
sys.exit(1)
114
+
115
+
try:
116
+
snapshot_download(**kwargs_dict)
117
+
except requests.exceptions.HTTPError as e:
118
+
status = getattr(e.response, "status_code", None)
119
+
url = getattr(e.response, "url", None)
120
+
if status == 401:
121
+
print(f"[ERROR] Unauthorized (401): Check your HuggingFace token or login.\nURL: {url}")
122
+
sys.exit(10)
123
+
elif status == 403:
124
+
print(f"[ERROR] Forbidden (403): Access denied.\nURL: {url}")
125
+
sys.exit(11)
126
+
elif status == 404:
127
+
print(f"[ERROR] Not Found (404): Resource does not exist.\nURL: {url}")
128
+
sys.exit(12)
129
+
else:
130
+
print(f"[ERROR] HTTP Error: {e}")
131
+
sys.exit(2)
132
+
except requests.exceptions.ConnectionError:
133
+
print("[ERROR] Network connection error. Please check your internet connection and try again.")
134
+
sys.exit(20)
135
+
except PermissionError as e:
136
+
print(f"[ERROR] Permission denied: {e.filename if hasattr(e, 'filename') else 'check file permissions'}")
137
+
print(" Ensure you have write access to the cache directory.")
138
+
sys.exit(13)
139
+
except OSError as e:
140
+
import errno
141
+
if e.errno == errno.ENOSPC:
142
+
print("[ERROR] No space left on device. Please free up disk space and try again.")
143
+
sys.exit(14)
144
+
elif e.errno == errno.EACCES:
145
+
print(f"[ERROR] Access denied: {e.filename if hasattr(e, 'filename') else 'check permissions'}")
146
+
sys.exit(13)
147
+
else:
148
+
print(f"[ERROR] OS Error during download: {e}")
149
+
sys.exit(15)
150
+
except Exception as e:
151
+
print(f"[ERROR] Unexpected error during download: {type(e).__name__}: {e}")
152
+
sys.exit(2)
153
+
finally:
154
+
try:
155
+
os.unlink(kwargs_file)
156
+
except Exception:
157
+
pass
158
+
159
+
sys.exit(0)
160
+
161
+
if __name__ == "__main__":
162
+
main()
163
+
+2
server/uv.lock
+2
server/uv.lock
···
980
980
dependencies = [
981
981
{ name = "black" },
982
982
{ name = "fastapi" },
983
+
{ name = "huggingface-hub" },
983
984
{ name = "mlx-lm" },
984
985
{ name = "uvicorn" },
985
986
]
···
988
989
requires-dist = [
989
990
{ name = "black" },
990
991
{ name = "fastapi" },
992
+
{ name = "huggingface-hub", specifier = ">=0.34.0" },
991
993
{ name = "mlx-lm" },
992
994
{ name = "uvicorn" },
993
995
]
+24
-2
src/runner/mlx.rs
+24
-2
src/runner/mlx.rs
···
2
2
use anyhow::{Context, Result};
3
3
use futures_util::StreamExt;
4
4
use owo_colors::OwoColorize;
5
-
use reqwest::Client;
5
+
use reqwest::{Client, StatusCode};
6
6
use serde_json::{Value, json};
7
7
use std::io::Write;
8
8
use std::path::PathBuf;
···
201
201
.send()
202
202
.await
203
203
.unwrap();
204
+
match res.status() {
205
+
StatusCode::OK => Ok(()),
206
+
StatusCode::NOT_FOUND => download_model(model_name).await,
207
+
_ => {
208
+
println!("err {:?}", res);
209
+
Ok(())
210
+
}
211
+
}
212
+
}
213
+
214
+
async fn download_model(model_name: &str) -> Result<(), String> {
215
+
println!("Downloading the model {} ....", model_name);
216
+
let client = Client::new();
217
+
let body = json!({
218
+
"model": model_name
219
+
});
220
+
let res = client
221
+
.post("http://127.0.0.1:6969/download")
222
+
.json(&body)
223
+
.send()
224
+
.await
225
+
.unwrap();
204
226
if res.status() == 200 {
205
227
Ok(())
206
228
} else {
207
-
Err(String::from("request failed"))
229
+
Err(String::from("Downloading model failed"))
208
230
}
209
231
}
210
232