@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.
- package/gonext-local-worker.mjs +53 -0
- package/gonext_agent_chat.py +147 -44
- package/gonext_mlx_embed.py +155 -0
- package/package.json +2 -1
package/gonext-local-worker.mjs
CHANGED
|
@@ -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
|
package/gonext_agent_chat.py
CHANGED
|
@@ -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
|
|
1269
|
-
"""Embed a batch of texts
|
|
1270
|
-
|
|
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
|
|
1275
|
+
raise RuntimeError("No embedding endpoint configured (set the RAG embedding server URL).")
|
|
1273
1276
|
ctx = _ssl_context()
|
|
1274
|
-
|
|
1275
|
-
|
|
1277
|
+
last_err = None
|
|
1278
|
+
|
|
1279
|
+
def _post(url, payload, timeout):
|
|
1276
1280
|
req = urllib.request.Request(
|
|
1277
|
-
|
|
1278
|
-
|
|
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=
|
|
1283
|
-
|
|
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
|
|
1289
|
-
|
|
1290
|
-
|
|
1291
|
-
|
|
1292
|
-
|
|
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
|
-
|
|
1296
|
-
|
|
1297
|
-
f"{root}/api/embeddings",
|
|
1298
|
-
data
|
|
1299
|
-
|
|
1300
|
-
|
|
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
|
-
|
|
1303
|
-
|
|
1304
|
-
|
|
1305
|
-
|
|
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
|
-
#
|
|
1518
|
-
#
|
|
1519
|
-
|
|
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
|
|
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 +=
|
|
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
|
|
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
|
-
" -
|
|
1719
|
-
"
|
|
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
|
|
1725
|
-
"
|
|
1726
|
-
"
|
|
1727
|
-
"
|
|
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 =
|
|
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 =
|
|
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 =
|
|
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,
|
|
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.
|
|
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
|
],
|