@tiens.nguyen/gonext-local-worker 1.0.157 → 1.0.160

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -5,6 +5,7 @@
5
5
  * writes ~/.gonext/worker.env
6
6
  * - `gonext-local-worker ws-ping-test` — POST worker ws-ping (needs GONEXT_* env)
7
7
  * - `gonext-local-worker simulate-chat [text]` — claim next chat job, push fake reply like the real worker (needs GONEXT_* env)
8
+ * - `gonext-local-worker embed --model <mlx-embed-model> [--port 8085]` — run the MLX embeddings server (RAG); use a separate terminal
8
9
  * - `gonext-local-worker` — starts polling loop (claims jobs and runs models)
9
10
  */
10
11
  import { mkdir, mkdtemp, readFile, rm, stat, writeFile } from "node:fs/promises";
@@ -221,6 +222,57 @@ if (args[0] === "ws-ping-test") {
221
222
  process.exit(0);
222
223
  }
223
224
 
225
+ if (args[0] === "embed") {
226
+ // Start the MLX embeddings server (OpenAI-compatible /v1/embeddings) for the RAG
227
+ // feature, in the FOREGROUND of this terminal. Run it in a SEPARATE terminal from the
228
+ // worker (the worker keeps polling; this holds one terminal for the embed model).
229
+ // gonext-local-worker embed --model ~/mlx-models/<embed-model> [--port 8085] [--host 127.0.0.1]
230
+ // Then set web Settings → Agent → RAG → "Embedding server URL" to http://<host>:<port>.
231
+ const flag = (name, fallback) => {
232
+ const i = args.indexOf(name);
233
+ return i >= 0 && typeof args[i + 1] === "string" ? args[i + 1] : fallback;
234
+ };
235
+ let model = flag("--model", process.env.GONEXT_EMBED_MODEL || "");
236
+ const port = flag("--port", process.env.GONEXT_EMBED_PORT || "8085");
237
+ const host = flag("--host", process.env.GONEXT_EMBED_HOST || "127.0.0.1");
238
+ if (!model) {
239
+ console.error(
240
+ "[gonext-worker] embed: --model is required.\n" +
241
+ " Usage: gonext-local-worker embed --model ~/mlx-models/Qwen3-Embedding-8B-4bit-DWQ [--port 8085] [--host 127.0.0.1]\n" +
242
+ " (or set GONEXT_EMBED_MODEL). Then set the RAG embedding server URL to http://<host>:<port> in web Settings."
243
+ );
244
+ process.exit(1);
245
+ }
246
+ // Expand a leading ~ to the home dir (spawn does not do shell expansion).
247
+ if (model.startsWith("~/") || model === "~") {
248
+ model = join(homedir(), model.slice(model.length > 1 ? 2 : 1));
249
+ }
250
+ const python =
251
+ (process.env.GONEXT_PROBE_PYTHON ?? process.env.GONEXT_MLX_LM_PYTHON ?? "").trim() ||
252
+ "python3";
253
+ const script = join(WORKER_DIR, "gonext_mlx_embed.py");
254
+ console.log(
255
+ `[gonext-worker] starting MLX embeddings server → ${python} ${script} --model ${model} --port ${port} --host ${host}`
256
+ );
257
+ const child = spawn(
258
+ python,
259
+ [script, "--model", model, "--port", String(port), "--host", String(host)],
260
+ { stdio: "inherit" }
261
+ );
262
+ child.on("error", (err) => {
263
+ console.error(`[gonext-worker] embed: failed to start python (${python}): ${err.message}`);
264
+ process.exit(1);
265
+ });
266
+ // Forward Ctrl-C to the child so the model unloads cleanly.
267
+ const stop = () => child.kill("SIGINT");
268
+ process.on("SIGINT", stop);
269
+ process.on("SIGTERM", stop);
270
+ // Block the module from continuing to the polling loop; exit when the child exits.
271
+ await new Promise((resolve) => child.on("exit", (code) => resolve(code ?? 0))).then(
272
+ (code) => process.exit(code)
273
+ );
274
+ }
275
+
224
276
  if (args[0] === "simulate-chat") {
225
277
  assertWorkerRuntimeConfig("simulate-chat");
226
278
  const reply =
@@ -1610,6 +1662,7 @@ async function runAgentChatJob(job) {
1610
1662
  ragAwsAccessKeyId: payload?.ragAwsAccessKeyId ?? "",
1611
1663
  ragAwsSecretAccessKey: payload?.ragAwsSecretAccessKey ?? "",
1612
1664
  ragEmbedModel: payload?.ragEmbedModel ?? "",
1665
+ ragEmbedUrl: payload?.ragEmbedUrl ?? "",
1613
1666
  ragTopK: payload?.ragTopK ?? 6,
1614
1667
  });
1615
1668
  // 30 min max for an agent run: multi-step ReAct on a 14B/31B with a cold prompt
@@ -1265,44 +1265,70 @@ def _rag_chunk_text(text: str, rel_path: str, ext: str):
1265
1265
  return chunks
1266
1266
 
1267
1267
 
1268
- def _ollama_embed(ollama_root: str, model: str, texts: list) -> list:
1269
- """Embed a batch of texts via Ollama /api/embed (falls back to /api/embeddings)."""
1270
- root = (ollama_root or "").rstrip("/")
1268
+ def _embed(base: str, model: str, texts: list) -> list:
1269
+ """Embed a batch of texts. Auto-detects the endpoint from `base` (a bare host:port,
1270
+ an .../v1, or an Ollama host): tries OpenAI-compatible /v1/embeddings first (MLX
1271
+ servers, modern Ollama), then Ollama /api/embed, then legacy /api/embeddings."""
1272
+ # Normalize to a bare root (strip any /v1, /api, /api/embed(dings) suffix).
1273
+ root = re.sub(r"/(v1|api/embeddings|api/embed|api)/?$", "", (base or "").rstrip("/")).rstrip("/")
1271
1274
  if not root:
1272
- raise RuntimeError("No Ollama endpoint available for embeddings.")
1275
+ raise RuntimeError("No embedding endpoint configured (set the RAG embedding server URL).")
1273
1276
  ctx = _ssl_context()
1274
- # Preferred: batch /api/embed (recent Ollama).
1275
- try:
1277
+ last_err = None
1278
+
1279
+ def _post(url, payload, timeout):
1276
1280
  req = urllib.request.Request(
1277
- f"{root}/api/embed",
1278
- data=json.dumps({"model": model, "input": texts}).encode("utf-8"),
1279
- headers={"Content-Type": "application/json"},
1280
- method="POST",
1281
+ url, data=json.dumps(payload).encode("utf-8"),
1282
+ headers={"Content-Type": "application/json"}, method="POST",
1281
1283
  )
1282
- with urllib.request.urlopen(req, timeout=300, context=ctx) as resp:
1283
- data = json.loads(resp.read().decode("utf-8", "replace"))
1284
+ with urllib.request.urlopen(req, timeout=timeout, context=ctx) as resp:
1285
+ return json.loads(resp.read().decode("utf-8", "replace"))
1286
+
1287
+ # 1) OpenAI-compatible /v1/embeddings — MLX servers on a port, and modern Ollama.
1288
+ try:
1289
+ data = _post(f"{root}/v1/embeddings", {"model": model, "input": texts}, 300)
1290
+ items = data.get("data")
1291
+ if isinstance(items, list) and items and isinstance(items[0], dict) and "embedding" in items[0]:
1292
+ return [it["embedding"] for it in items]
1293
+ except urllib.error.HTTPError as e:
1294
+ if e.code not in (400, 404, 405, 501):
1295
+ raise RuntimeError(f"Embeddings failed (HTTP {e.code}): {e.read().decode('utf-8', 'replace')[:200]}")
1296
+ last_err = f"/v1/embeddings HTTP {e.code}"
1297
+ except urllib.error.URLError as e:
1298
+ raise RuntimeError(f"Embedding server unreachable at {root}: {getattr(e, 'reason', e)}")
1299
+
1300
+ # 2) Ollama batch /api/embed.
1301
+ try:
1302
+ data = _post(f"{root}/api/embed", {"model": model, "input": texts}, 300)
1284
1303
  embs = data.get("embeddings")
1285
1304
  if isinstance(embs, list) and embs and isinstance(embs[0], list):
1286
1305
  return embs
1287
1306
  except urllib.error.HTTPError as e:
1288
- if e.code != 404:
1289
- detail = e.read().decode("utf-8", "replace")[:200]
1290
- raise RuntimeError(f"Ollama embed failed (HTTP {e.code}): {detail}")
1291
- except urllib.error.URLError as e:
1292
- raise RuntimeError(f"Ollama embed unreachable: {getattr(e, 'reason', e)}")
1293
- # Fallback: per-text /api/embeddings (older Ollama).
1307
+ if e.code not in (400, 404, 405, 501):
1308
+ raise RuntimeError(f"Ollama embed failed (HTTP {e.code}): {e.read().decode('utf-8', 'replace')[:200]}")
1309
+ last_err = f"/api/embed HTTP {e.code}"
1310
+
1311
+ # 3) Legacy Ollama /api/embeddings (per text).
1294
1312
  out = []
1295
- for t in texts:
1296
- req = urllib.request.Request(
1297
- f"{root}/api/embeddings",
1298
- data=json.dumps({"model": model, "prompt": t}).encode("utf-8"),
1299
- headers={"Content-Type": "application/json"},
1300
- method="POST",
1313
+ try:
1314
+ for t in texts:
1315
+ data = _post(f"{root}/api/embeddings", {"model": model, "prompt": t}, 120)
1316
+ out.append(data.get("embedding") or [])
1317
+ except (urllib.error.HTTPError, urllib.error.URLError) as e:
1318
+ code = getattr(e, "code", None)
1319
+ raise RuntimeError(
1320
+ f"No embeddings endpoint at {root} (tried /v1/embeddings, /api/embed, /api/embeddings; "
1321
+ f"last HTTP {code if code else getattr(e, 'reason', e)}). Make sure an EMBEDDINGS-capable "
1322
+ f"server is running there and the model '{model}' is available "
1323
+ f"(e.g. `ollama pull {model}`). NOTE: mlx_lm.server does NOT serve embeddings — point the "
1324
+ f"RAG embedding server URL at Ollama (or another embeddings server) instead."
1301
1325
  )
1302
- with urllib.request.urlopen(req, timeout=120, context=ctx) as resp:
1303
- data = json.loads(resp.read().decode("utf-8", "replace"))
1304
- out.append(data.get("embedding") or [])
1305
- return out
1326
+ if any(any(v) for v in out):
1327
+ return out
1328
+ raise RuntimeError(
1329
+ f"Embeddings server at {root} returned empty vectors for model '{model}' "
1330
+ f"(is '{model}' an embedding model that is installed there?)."
1331
+ )
1306
1332
 
1307
1333
 
1308
1334
  def _cosine(a: list, b: list) -> float:
@@ -1336,6 +1362,24 @@ def _rag_parse_location(loc: str):
1336
1362
  return bucket, prefix.strip("/")
1337
1363
 
1338
1364
 
1365
+ def _rag_read_allowed(path: str) -> str:
1366
+ """Resolve `path` and ensure it stays within a safe root (worker cwd, temp dir, or
1367
+ ~/.gonext) — so the agent's file tools can't read arbitrary files like ~/.ssh."""
1368
+ import os
1369
+ import tempfile
1370
+ rp = os.path.realpath(path)
1371
+ roots = [
1372
+ os.path.realpath(os.getcwd()),
1373
+ os.path.realpath(tempfile.gettempdir()),
1374
+ os.path.realpath(os.path.join(os.path.expanduser("~"), ".gonext")),
1375
+ ]
1376
+ if any(rp == r or rp.startswith(r + os.sep) for r in roots):
1377
+ return rp
1378
+ raise ValueError(
1379
+ "Path is outside the allowed working directories (download/unzip outputs only)."
1380
+ )
1381
+
1382
+
1339
1383
  def _rag_s3_client(region: str, akid: str, secret: str):
1340
1384
  import boto3 # noqa: PLC0415
1341
1385
  return boto3.client(
@@ -1509,20 +1553,22 @@ def run_agent_chat(cfg):
1509
1553
  rag_akid = (cfg.get("ragAwsAccessKeyId") or "").strip()
1510
1554
  rag_secret = (cfg.get("ragAwsSecretAccessKey") or "").strip()
1511
1555
  rag_embed_model = (cfg.get("ragEmbedModel") or "nomic-embed-text").strip()
1556
+ rag_embed_url = (cfg.get("ragEmbedUrl") or "").strip()
1512
1557
  try:
1513
1558
  rag_top_k = int(cfg.get("ragTopK") or 6)
1514
1559
  except (TypeError, ValueError):
1515
1560
  rag_top_k = 6
1516
1561
  rag_top_k = max(1, min(50, rag_top_k))
1517
- # Embeddings run on the Ollama coding server (same box as the coder). Strip the
1518
- # OpenAI-compat /v1 suffix to reach the native /api/embed root.
1519
- rag_ollama_root = re.sub(r"/v1/?$", "", (coding_base_url or agent_base_url or "").rstrip("/"))
1562
+ # Embedding server: an explicit MLX/Ollama endpoint from Settings if given, else the
1563
+ # agent's Ollama coding server. _embed() auto-detects OpenAI-compat (/v1/embeddings,
1564
+ # e.g. an MLX server on a port) vs Ollama (/api/embed) — so a bare host:port works.
1565
+ rag_embed_base = rag_embed_url or (coding_base_url or agent_base_url or "")
1520
1566
  try:
1521
1567
  import boto3 as _boto3_probe # noqa: F401,PLC0415
1522
1568
  _boto3_ok = True
1523
1569
  except Exception: # noqa: BLE001
1524
1570
  _boto3_ok = False
1525
- _RAG_AVAILABLE = bool(rag_enabled and rag_akid and rag_secret and rag_ollama_root and _boto3_ok)
1571
+ _RAG_AVAILABLE = bool(rag_enabled and rag_akid and rag_secret and rag_embed_base and _boto3_ok)
1526
1572
  if rag_enabled and not _boto3_ok:
1527
1573
  _log("RAG requested but boto3 is not installed in the worker python — RAG tools disabled")
1528
1574
 
@@ -1709,22 +1755,29 @@ def run_agent_chat(cfg):
1709
1755
  _tool_num += 1
1710
1756
  _email_num = _tool_num
1711
1757
  if _RAG_AVAILABLE:
1712
- _tool_num += 5 # download_file, unzip_file, rag_index, rag_add, rag_search
1758
+ _tool_num += 7 # download_file, unzip_file, list_dir, read_text_file, rag_index/add/search
1713
1759
  _tool_count_word = {6: "SIX", 7: "SEVEN", 8: "EIGHT"}.get(_tool_num, str(_tool_num))
1714
1760
  _rag_tool_block = (
1715
- " KNOWLEDGE BASE (RAG) — to SUMMARIZE or ANSWER questions about a ZIP of files at a URL:\n"
1761
+ " FILES & KNOWLEDGE BASE — to SUMMARIZE or ANSWER questions about a ZIP of files at a URL:\n"
1716
1762
  " - download_file(url) — download a file (e.g. a .zip) to this machine.\n"
1717
1763
  " - unzip_file(zip_path) — extract a downloaded .zip locally.\n"
1718
- " - rag_index(path, source_url) — index the unzipped text/code files into the knowledge "
1719
- "base (embeds to S3). source_url MUST be the ORIGINAL url the user gave.\n"
1764
+ " - list_dir(path) — list the files inside an unzipped folder.\n"
1765
+ " - read_text_file(path) read ONE file's text (e.g. a README) directly. Best for a QUICK "
1766
+ "project summary — you do NOT need to index for that.\n"
1767
+ " - rag_index(path, source_url) — index MANY text/code files into a searchable knowledge "
1768
+ "base (embeds to S3). source_url MUST be the ORIGINAL url the user gave. Use for deep Q&A "
1769
+ "across the whole project.\n"
1720
1770
  " - rag_add(source_url, text) — add extra info the user provides to that knowledge base.\n"
1721
1771
  " - rag_search(source_url, query) — retrieve the most relevant chunks to answer a question.\n"
1772
+ " NOTE: you may ONLY read files with these tools — plain `import os`/`open()` is blocked.\n"
1722
1773
  ) if _RAG_AVAILABLE else ""
1723
1774
  _rag_choose_line = (
1724
- "- user gives a URL to a ZIP of files and asks to summarize/analyze/RAG them -> "
1725
- "download_file(url) THEN unzip_file(zip) THEN rag_index(dir, source_url=url) THEN "
1726
- "rag_search(source_url=url, query=<the user's question>), THEN answer from the results. "
1727
- "For FOLLOW-UP questions about the SAME url, just rag_search (do NOT re-download). "
1775
+ "- user gives a URL to a ZIP and asks to SUMMARIZE it -> download_file(url) THEN unzip_file(zip) "
1776
+ "THEN list_dir(dir) to find the README, THEN read_text_file(that README) THEN answer from it. "
1777
+ "This is the FASTEST path for a summary do NOT index unless deep Q&A across many files is needed.\n"
1778
+ "- user wants deep Q&A across the WHOLE project -> download_file THEN unzip_file THEN "
1779
+ "rag_index(dir, source_url=url) THEN rag_search(source_url=url, query=<question>) THEN answer. "
1780
+ "For FOLLOW-UPs on the SAME url, just rag_search (do NOT re-download). "
1728
1781
  "If the user adds information, rag_add(source_url=url, text=...).\n"
1729
1782
  ) if _RAG_AVAILABLE else ""
1730
1783
  _pdf_read_tool_line = (
@@ -2650,6 +2703,55 @@ def run_agent_chat(cfg):
2650
2703
  except Exception as e: # noqa: BLE001
2651
2704
  return f"Error: unzip failed: {type(e).__name__}: {e}"
2652
2705
 
2706
+ @tool
2707
+ def list_dir(path: str) -> str:
2708
+ """List files and subfolders under a local directory (from unzip_file). Use to see what files are in the project before reading or indexing them.
2709
+
2710
+ Args:
2711
+ path: local directory path to list.
2712
+ """
2713
+ try:
2714
+ import os
2715
+ root = _rag_read_allowed((path or ".").strip())
2716
+ if os.path.isfile(root):
2717
+ return f"{root} is a file ({os.path.getsize(root)} bytes)."
2718
+ entries = []
2719
+ for dirpath, dirnames, filenames in os.walk(root):
2720
+ dirnames[:] = [d for d in dirnames if d not in _RAG_SKIP_DIRS]
2721
+ for fn in filenames:
2722
+ rel = os.path.relpath(os.path.join(dirpath, fn), root)
2723
+ entries.append(rel)
2724
+ if len(entries) >= 300:
2725
+ break
2726
+ if len(entries) >= 300:
2727
+ break
2728
+ entries.sort()
2729
+ more = " …(truncated)" if len(entries) >= 300 else ""
2730
+ return f"{len(entries)} files under {root}:\n" + "\n".join(entries) + more
2731
+ except Exception as e: # noqa: BLE001
2732
+ return f"Error: list_dir failed: {type(e).__name__}: {e}"
2733
+
2734
+ @tool
2735
+ def read_text_file(path: str, max_chars: int = 20000) -> str:
2736
+ """Read the text contents of a single local file (e.g. a README) from a downloaded/unzipped project. Use this for a QUICK summary when you don't need full indexing.
2737
+
2738
+ Args:
2739
+ path: local path to the text file (e.g. project_files/tools-hook/README.md).
2740
+ max_chars: maximum characters to return (default 20000).
2741
+ """
2742
+ try:
2743
+ import os
2744
+ rp = _rag_read_allowed((path or "").strip())
2745
+ if not os.path.isfile(rp):
2746
+ return f"Error: not a file: {path}"
2747
+ if os.path.getsize(rp) > 5 * 1024 * 1024:
2748
+ return "Error: file too large to read (> 5 MB)."
2749
+ with open(rp, "r", encoding="utf-8", errors="replace") as fh:
2750
+ data = fh.read(max(500, min(int(max_chars or 20000), 100000)))
2751
+ return data or "(empty file)"
2752
+ except Exception as e: # noqa: BLE001
2753
+ return f"Error: read_text_file failed: {type(e).__name__}: {e}"
2754
+
2653
2755
  @tool
2654
2756
  def rag_index(path: str, source_url: str) -> str:
2655
2757
  """Index a local file or directory of TEXT/code files into the knowledge base for source_url (stores embeddings on S3). Use after unzip_file so the files become searchable. source_url MUST be the original URL the user provided.
@@ -2678,7 +2780,7 @@ def run_agent_chat(cfg):
2678
2780
  return "No indexable text files found at that path."
2679
2781
  for i in range(0, len(records), 64):
2680
2782
  batch = records[i:i + 64]
2681
- vecs = _ollama_embed(rag_ollama_root, rag_embed_model, [r["text"] for r in batch])
2783
+ vecs = _embed(rag_embed_base, rag_embed_model, [r["text"] for r in batch])
2682
2784
  for r, v in zip(batch, vecs):
2683
2785
  r["embedding"] = v
2684
2786
  _emit({"type": "step", "text": f"Embedded {min(i + 64, len(records))}/{len(records)} chunks…"})
@@ -2718,7 +2820,7 @@ def run_agent_chat(cfg):
2718
2820
  try:
2719
2821
  import time as _t
2720
2822
  records = _rag_chunk_text(text, "user-note", ".txt")
2721
- vecs = _ollama_embed(rag_ollama_root, rag_embed_model, [r["text"] for r in records])
2823
+ vecs = _embed(rag_embed_base, rag_embed_model, [r["text"] for r in records])
2722
2824
  for r, v in zip(records, vecs):
2723
2825
  r["embedding"] = v
2724
2826
  client, bucket, prefix = _rag_s3_and_loc()
@@ -2749,7 +2851,7 @@ def run_agent_chat(cfg):
2749
2851
  chunks = _rag_load_chunks(client, bucket, base)
2750
2852
  if not chunks:
2751
2853
  return f"No knowledge base found for {source_url}. Run rag_index first."
2752
- qvec = _ollama_embed(rag_ollama_root, rag_embed_model, [query])[0]
2854
+ qvec = _embed(rag_embed_base, rag_embed_model, [query])[0]
2753
2855
  scored = sorted(
2754
2856
  ((_cosine(qvec, c.get("embedding") or []), c) for c in chunks),
2755
2857
  key=lambda x: x[0], reverse=True,
@@ -2770,7 +2872,8 @@ def run_agent_chat(cfg):
2770
2872
  if _EMAIL_AVAILABLE:
2771
2873
  agent_tools.append(send_email)
2772
2874
  if _RAG_AVAILABLE:
2773
- agent_tools += [download_file, unzip_file, rag_index, rag_add, rag_search]
2875
+ agent_tools += [download_file, unzip_file, list_dir, read_text_file,
2876
+ rag_index, rag_add, rag_search]
2774
2877
  agent_kwargs = dict(
2775
2878
  tools=agent_tools,
2776
2879
  model=model,
@@ -0,0 +1,155 @@
1
+ #!/usr/bin/env python3
2
+ """
3
+ GoNext MLX embeddings server — an OpenAI-compatible /v1/embeddings endpoint for MLX
4
+ embedding models (e.g. Qwen3-Embedding-8B), which `mlx_lm.server` does NOT provide.
5
+
6
+ The GoNext worker's RAG tools call `POST {url}/v1/embeddings` first, so pointing the
7
+ "RAG embedding server URL" (web Settings → Agent → RAG) at this server makes MLX
8
+ embeddings work with no other changes.
9
+
10
+ Run:
11
+ python3 gonext_mlx_embed.py --model ~/mlx-models/Qwen3-Embedding-8B-4bit-DWQ --port 8085
12
+
13
+ Then set the RAG embedding server URL to http://127.0.0.1:8085
14
+
15
+ Endpoints:
16
+ POST /v1/embeddings {"model": <ignored>, "input": "text" | ["t1","t2",...]}
17
+ → {"object":"list","data":[{"object":"embedding","index":i,
18
+ "embedding":[...]}], "model": <name>, "usage": {...}}
19
+ GET /v1/models → the single loaded model
20
+ GET /health → {"status":"ok"}
21
+
22
+ Pooling: last-token hidden state (Qwen3-Embedding convention), L2-normalized.
23
+ Requires: mlx, mlx_lm (already used by the worker). No extra deps.
24
+ """
25
+ import argparse
26
+ import json
27
+ import os
28
+ import sys
29
+ import threading
30
+ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
31
+
32
+ import mlx.core as mx
33
+ from mlx_lm import load
34
+
35
+ # Cap tokens per input so a pathologically long chunk can't OOM / stall the GPU.
36
+ MAX_TOKENS = 8192
37
+
38
+ _model = None
39
+ _tok = None
40
+ _inner = None # the transformer body that returns hidden states
41
+ _model_name = "mlx-embedding"
42
+ _lock = threading.Lock() # MLX eval isn't guaranteed thread-safe; serialize inference
43
+
44
+
45
+ def _load(model_path: str) -> None:
46
+ global _model, _tok, _inner, _model_name
47
+ print(f"[gonext-mlx-embed] loading {model_path} …", flush=True)
48
+ _model, _tok = load(model_path)
49
+ # mlx_lm wraps the transformer body as `.model` (returns hidden states); `_model(x)`
50
+ # would return logits from the LM head, which we do NOT want.
51
+ _inner = getattr(_model, "model", None) or getattr(_model, "language_model", None)
52
+ if _inner is None:
53
+ raise RuntimeError("Could not find the transformer body on the loaded model.")
54
+ _model_name = os.path.basename(os.path.normpath(model_path))
55
+ # Warm + sanity-check the embedding path so a bad model fails fast at startup.
56
+ v = _embed_one("warmup")
57
+ print(f"[gonext-mlx-embed] ready — {_model_name}, dim={len(v)}", flush=True)
58
+
59
+
60
+ def _embed_one(text: str) -> list:
61
+ ids = _tok.encode(text or " ")
62
+ if not ids:
63
+ ids = _tok.encode(" ")
64
+ if len(ids) > MAX_TOKENS:
65
+ ids = ids[:MAX_TOKENS]
66
+ h = _inner(mx.array([ids])) # (1, seq, hidden)
67
+ v = h[0, -1, :] # last-token pooling
68
+ v = v / mx.maximum(mx.linalg.norm(v), 1e-9) # L2 normalize
69
+ mx.eval(v)
70
+ return [float(x) for x in v.tolist()]
71
+
72
+
73
+ def _embed_batch(texts: list) -> list:
74
+ with _lock:
75
+ out = [_embed_one(t) for t in texts]
76
+ return out
77
+
78
+
79
+ class Handler(BaseHTTPRequestHandler):
80
+ protocol_version = "HTTP/1.1"
81
+
82
+ def log_message(self, *a): # quieter logs
83
+ pass
84
+
85
+ def _send(self, code: int, obj: dict) -> None:
86
+ body = json.dumps(obj).encode("utf-8")
87
+ self.send_response(code)
88
+ self.send_header("Content-Type", "application/json")
89
+ self.send_header("Content-Length", str(len(body)))
90
+ self.end_headers()
91
+ self.wfile.write(body)
92
+
93
+ def do_GET(self):
94
+ if self.path == "/health":
95
+ self._send(200, {"status": "ok", "model": _model_name})
96
+ elif self.path in ("/v1/models", "/models"):
97
+ self._send(200, {"object": "list", "data": [
98
+ {"id": _model_name, "object": "model", "owned_by": "gonext-mlx-embed"}
99
+ ]})
100
+ else:
101
+ self._send(404, {"error": {"message": f"unknown path {self.path}"}})
102
+
103
+ def do_POST(self):
104
+ if self.path not in ("/v1/embeddings", "/embeddings"):
105
+ self._send(404, {"error": {"message": f"unknown path {self.path}"}})
106
+ return
107
+ try:
108
+ n = int(self.headers.get("Content-Length", 0) or 0)
109
+ payload = json.loads(self.rfile.read(n) or b"{}")
110
+ except Exception as e: # noqa: BLE001
111
+ self._send(400, {"error": {"message": f"bad JSON: {e}"}})
112
+ return
113
+ inp = payload.get("input")
114
+ if inp is None:
115
+ self._send(400, {"error": {"message": "missing 'input'"}})
116
+ return
117
+ texts = [inp] if isinstance(inp, str) else list(inp)
118
+ if not all(isinstance(t, str) for t in texts):
119
+ self._send(400, {"error": {"message": "'input' must be a string or list of strings"}})
120
+ return
121
+ try:
122
+ vecs = _embed_batch(texts)
123
+ except Exception as e: # noqa: BLE001
124
+ self._send(500, {"error": {"message": f"embedding failed: {type(e).__name__}: {e}"}})
125
+ return
126
+ total_tokens = sum(len(_tok.encode(t or " ")) for t in texts)
127
+ self._send(200, {
128
+ "object": "list",
129
+ "data": [
130
+ {"object": "embedding", "index": i, "embedding": v}
131
+ for i, v in enumerate(vecs)
132
+ ],
133
+ "model": payload.get("model") or _model_name,
134
+ "usage": {"prompt_tokens": total_tokens, "total_tokens": total_tokens},
135
+ })
136
+
137
+
138
+ def main():
139
+ ap = argparse.ArgumentParser(description="OpenAI-compatible MLX embeddings server.")
140
+ ap.add_argument("--model", required=True, help="Path to the MLX embedding model dir.")
141
+ ap.add_argument("--port", type=int, default=8085)
142
+ ap.add_argument("--host", default="127.0.0.1")
143
+ args = ap.parse_args()
144
+ _load(os.path.expanduser(args.model))
145
+ srv = ThreadingHTTPServer((args.host, args.port), Handler)
146
+ print(f"[gonext-mlx-embed] serving /v1/embeddings on http://{args.host}:{args.port}", flush=True)
147
+ try:
148
+ srv.serve_forever()
149
+ except KeyboardInterrupt:
150
+ print("\n[gonext-mlx-embed] shutting down", flush=True)
151
+ srv.shutdown()
152
+
153
+
154
+ if __name__ == "__main__":
155
+ sys.exit(main())
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@tiens.nguyen/gonext-local-worker",
3
- "version": "1.0.157",
3
+ "version": "1.0.160",
4
4
  "description": "Polls GoNext cloud API for async local LLM jobs and runs them against Ollama/OpenAI-compatible servers on this Mac",
5
5
  "type": "module",
6
6
  "license": "MIT",
@@ -25,6 +25,7 @@
25
25
  "gonext_probe_agent.py",
26
26
  "gonext_agent_chat.py",
27
27
  "gonext_transcribe.py",
28
+ "gonext_mlx_embed.py",
28
29
  "README.md",
29
30
  "launchd/"
30
31
  ],