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

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.
@@ -143,6 +143,8 @@ const OCR_CORRECT_TIMEOUT_MS = 120_000;
143
143
  // cold load of a big/multimodal model into VRAM can take ~20s+; the request isn't
144
144
  // awaited, so a long ceiling just avoids aborting an in-progress load.
145
145
  const AGENT_WARMUP_TIMEOUT_MS = 180_000;
146
+ // How often the background keep-warm re-asserts the coding model's keep_alive:-1.
147
+ const AGENT_WARMUP_INTERVAL_MS = 240_000; // 4 min — under Ollama's 5-min default unload
146
148
  // OCR translation/correction is hardcoded to a managed Ollama service + model.
147
149
  // (The user-configurable "OCR translation model" chooser was removed.)
148
150
  const OCR_TRANSLATE_OLLAMA_URL = "https://ollama1.gomarsic.cc";
@@ -1436,33 +1438,37 @@ async function runTranscribeJob(job) {
1436
1438
  }
1437
1439
  }
1438
1440
 
1439
- /**
1440
- * Fire-and-forget pre-load of the agent's Ollama coding model. Ollama evicts idle
1441
- * models (default OLLAMA_KEEP_ALIVE=5m), so the first agent step otherwise blocks
1442
- * ~20s+ while llama-server spins up and pages the weights into VRAM. The agent
1443
- * talks to Ollama over the OpenAI-compat /v1 endpoint, which IGNORES keep_alive
1444
- * so we hit the NATIVE /api/generate here (no prompt = load-only, returns as soon
1445
- * as the model is resident) with keep_alive:-1 to trigger the load AND pin it in
1446
- * memory. Kicked off the moment the job is claimed so the load overlaps python
1447
- * startup + prompt assembly. Best-effort: a non-Ollama coding server just 404s and
1448
- * the failure is swallowed — the job proceeds regardless.
1449
- */
1450
- function warmUpOllamaCodingModel(codingBaseURL, codingModelId) {
1451
- if (!codingBaseURL || !codingModelId) return;
1452
- // codingBaseURL is normalized to …/v1 for OpenAI-compat; strip it to reach the
1453
- // Ollama root where /api/generate lives.
1454
- const root = String(codingBaseURL).replace(/\/+$/, "").replace(/\/v1$/i, "");
1441
+ // —— Ollama coding-model keep-warm ——————————————————————————————————————————————
1442
+ // Ollama evicts idle models (default OLLAMA_KEEP_ALIVE=5m), so a cold agent step
1443
+ // blocks ~20s+ while llama-server pages the weights into VRAM. Instead of warming on
1444
+ // EVERY question (noisy, redundant), we remember the coding model from agent jobs and
1445
+ // re-assert a keep_alive:-1 load on a BACKGROUND INTERVAL so it stays resident between
1446
+ // questions. The agent talks to Ollama over /v1 (which ignores keep_alive), so we hit
1447
+ // the NATIVE /api/generate (no prompt = load-only, returns once the model is resident).
1448
+ // Best-effort throughout: a non-Ollama coding server just 404s and it's swallowed.
1449
+ let _warmCoding = null; // { root, model }
1450
+ let _warmTimer = null;
1451
+ let _warmLoggedOk = false;
1452
+
1453
+ function _warmupPing(quiet) {
1454
+ if (!_warmCoding) return;
1455
+ const { root, model } = _warmCoding;
1455
1456
  const url = `${root}/api/generate`;
1456
1457
  fetch(url, {
1457
1458
  method: "POST",
1458
1459
  headers: { "Content-Type": "application/json" },
1459
- body: JSON.stringify({ model: codingModelId, keep_alive: -1 }),
1460
+ body: JSON.stringify({ model, keep_alive: -1 }),
1460
1461
  signal: AbortSignal.timeout(AGENT_WARMUP_TIMEOUT_MS),
1461
1462
  })
1462
1463
  .then((res) => {
1463
- console.log(
1464
- `[gonext-worker] agent warm-up ${res.ok ? "ok (model loading/resident)" : `HTTP ${res.status}`} → ${url} model=${codingModelId}`
1465
- );
1464
+ if (!res.ok) {
1465
+ console.log(`[gonext-worker] agent warm-up HTTP ${res.status} → ${url} model=${model}`);
1466
+ } else if (!quiet || !_warmLoggedOk) {
1467
+ // Quiet interval pings only log the first success per model; a fresh/changed
1468
+ // model always logs so you can see it warm.
1469
+ console.log(`[gonext-worker] agent warm-up ok (model loading/resident) → ${url} model=${model}`);
1470
+ _warmLoggedOk = true;
1471
+ }
1466
1472
  })
1467
1473
  .catch((err) => {
1468
1474
  const reason =
@@ -1471,12 +1477,37 @@ function warmUpOllamaCodingModel(codingBaseURL, codingModelId) {
1471
1477
  });
1472
1478
  }
1473
1479
 
1480
+ /**
1481
+ * Remember the agent's Ollama coding model (from a claimed agent_chat job) and keep it
1482
+ * resident via a background interval. Warms IMMEDIATELY only when the model is first
1483
+ * seen or CHANGES — not on every question — so repeat questions don't re-warm/re-log.
1484
+ */
1485
+ function trackCodingModelForWarmup(codingBaseURL, codingModelId) {
1486
+ if (!codingBaseURL || !codingModelId) return;
1487
+ // codingBaseURL is normalized to …/v1 for OpenAI-compat; strip it to reach the
1488
+ // Ollama root where /api/generate lives.
1489
+ const root = String(codingBaseURL).replace(/\/+$/, "").replace(/\/v1$/i, "");
1490
+ const changed = !_warmCoding || _warmCoding.root !== root || _warmCoding.model !== codingModelId;
1491
+ _warmCoding = { root, model: codingModelId };
1492
+ if (changed) {
1493
+ _warmLoggedOk = false;
1494
+ _warmupPing(false); // warm the newly-seen/changed model now so the next question isn't cold
1495
+ }
1496
+ if (!_warmTimer) {
1497
+ // 4 min — under Ollama's 5-min default unload, so the model stays resident even
1498
+ // if a proxy strips keep_alive. unref() so this timer never blocks worker exit.
1499
+ _warmTimer = setInterval(() => _warmupPing(true), AGENT_WARMUP_INTERVAL_MS);
1500
+ if (typeof _warmTimer.unref === "function") _warmTimer.unref();
1501
+ }
1502
+ }
1503
+
1474
1504
  async function runAgentChatJob(job) {
1475
1505
  const { jobId, payload } = job;
1476
1506
  const start = Date.now();
1477
- // Kick the Ollama coding model's cold load off immediately (non-blocking) so it
1478
- // overlaps prompt assembly instead of stalling the first step. See helper above.
1479
- warmUpOllamaCodingModel(payload?.codingBaseURL ?? "", payload?.codingModelId ?? "");
1507
+ // Remember this job's Ollama coding model and keep it resident via a background
1508
+ // interval. Only warms now if the model is new/changed repeat questions rely on
1509
+ // the interval, so we don't re-warm (or re-log) on every question.
1510
+ trackCodingModelForWarmup(payload?.codingBaseURL ?? "", payload?.codingModelId ?? "");
1480
1511
  const runRes = await workerFetch(`/api/worker/jobs/${jobId}`, {
1481
1512
  method: "PATCH",
1482
1513
  body: JSON.stringify({ jobStatus: "running" }),
@@ -1570,6 +1601,16 @@ async function runAgentChatJob(job) {
1570
1601
  emailBodyTemplate: payload?.emailBodyTemplate ?? "",
1571
1602
  emailFrom: payload?.emailFrom ?? "",
1572
1603
  emailAllowList: payload?.emailAllowList ?? "",
1604
+ // RAG tools (download/unzip/index/search over a user's ZIP-at-URL). Only enabled
1605
+ // + registered when ragEnabled + AWS creds present. The worker talks to the user's
1606
+ // OWN S3 bucket directly with these creds (boto3), NOT via presigned URLs.
1607
+ ragEnabled: payload?.ragEnabled ?? false,
1608
+ ragS3Location: payload?.ragS3Location ?? "",
1609
+ ragAwsRegion: payload?.ragAwsRegion ?? "",
1610
+ ragAwsAccessKeyId: payload?.ragAwsAccessKeyId ?? "",
1611
+ ragAwsSecretAccessKey: payload?.ragAwsSecretAccessKey ?? "",
1612
+ ragEmbedModel: payload?.ragEmbedModel ?? "",
1613
+ ragTopK: payload?.ragTopK ?? 6,
1573
1614
  });
1574
1615
  // 30 min max for an agent run: multi-step ReAct on a 14B/31B with a cold prompt
1575
1616
  // cache — or a remote Ollama box on slow GPU cold-loading a big model — can run
@@ -518,7 +518,7 @@ _AGENT_KEYWORDS = re.compile(
518
518
  r"|api|endpoint|url|http|https"
519
519
  r"|external\s+source|external\s+api|external\s+service"
520
520
  r"|web\s+service|rest\s+api|rest\s+call"
521
- r"|download|scrape|crawl"
521
+ r"|download|scrape|crawl|zip|unzip|\.zip|rag|index\s+the|knowledge\s+base|summari[sz]e"
522
522
  r"|search|find|look\s*up|lookup|weather|news|latest|current|today|tonight"
523
523
  r"|date|time|what\s+day|what\s+time"
524
524
  r"|read\s+(?:this\s+|that\s+|the\s+)?(?:page|article|url|link)|open\s+(?:this\s+|that\s+|the\s+)?(?:url|link|page)|contents?\s+of"
@@ -1080,6 +1080,270 @@ def _derive_pdf_title(doc_text: str) -> str:
1080
1080
  return "Document"
1081
1081
 
1082
1082
 
1083
+ # ============================ RAG (knowledge base) ============================
1084
+ # The agent can download a ZIP-at-URL, unzip it locally, index the text files into a
1085
+ # flat JSON vector store on the USER'S OWN S3 bucket (explicit AWS creds — NOT presigned),
1086
+ # then retrieve relevant chunks to answer questions. Embeddings run on the Ollama coding
1087
+ # server (/api/embed). Everything below is best-effort with hard safety caps.
1088
+
1089
+ _RAG_TEXT_EXTS = {
1090
+ ".cs", ".js", ".jsx", ".ts", ".tsx", ".mjs", ".cjs", ".vue", ".svelte",
1091
+ ".json", ".jsonc", ".md", ".mdx", ".txt", ".rst", ".py", ".rb", ".go",
1092
+ ".rs", ".java", ".kt", ".swift", ".php", ".c", ".h", ".cpp", ".hpp", ".cc",
1093
+ ".html", ".htm", ".css", ".scss", ".sass", ".less", ".xml", ".yaml", ".yml",
1094
+ ".toml", ".ini", ".cfg", ".conf", ".env", ".sh", ".bash", ".zsh", ".sql",
1095
+ ".graphql", ".proto", ".gradle", ".dockerfile", ".tsv", ".csv",
1096
+ }
1097
+ _RAG_SKIP_DIRS = {
1098
+ "node_modules", ".git", ".svn", ".hg", "__pycache__", ".venv", "venv",
1099
+ "dist", "build", ".next", ".nuxt", "out", ".cache", "bin", "obj", ".idea",
1100
+ ".gradle", "vendor", "target",
1101
+ }
1102
+ _RAG_MAX_DOWNLOAD_BYTES = 500 * 1024 * 1024 # 500 MB zip cap
1103
+ _RAG_MAX_UNZIP_BYTES = 2 * 1024 * 1024 * 1024 # 2 GB uncompressed (zip-bomb guard)
1104
+ _RAG_MAX_FILES = 20000
1105
+ _RAG_MAX_FILE_BYTES = 2 * 1024 * 1024 # skip individual text files > 2 MB
1106
+ _RAG_CHUNK_CHARS = 3200 # ~800 tokens
1107
+ _RAG_CHUNK_OVERLAP = 400
1108
+
1109
+
1110
+ def _rag_source_key(url: str) -> str:
1111
+ import hashlib
1112
+ return hashlib.sha256((url or "").strip().encode("utf-8")).hexdigest()[:32]
1113
+
1114
+
1115
+ def _rag_gdrive_direct(url: str) -> str:
1116
+ """Rewrite a Google Drive share link to its direct-download form."""
1117
+ m = re.search(r"drive\.google\.com/file/d/([A-Za-z0-9_-]+)", url or "")
1118
+ if m:
1119
+ return f"https://drive.google.com/uc?export=download&id={m.group(1)}"
1120
+ m = re.search(r"drive\.google\.com/open\?id=([A-Za-z0-9_-]+)", url or "")
1121
+ if m:
1122
+ return f"https://drive.google.com/uc?export=download&id={m.group(1)}"
1123
+ return url
1124
+
1125
+
1126
+ def _rag_assert_safe_url(url: str) -> None:
1127
+ """SSRF guard: only http/https, and refuse private/link-local/loopback hosts."""
1128
+ import ipaddress
1129
+ import socket
1130
+ from urllib.parse import urlparse
1131
+ p = urlparse(url or "")
1132
+ if p.scheme not in ("http", "https"):
1133
+ raise ValueError("Only http/https URLs are allowed.")
1134
+ host = p.hostname or ""
1135
+ if not host:
1136
+ raise ValueError("URL has no host.")
1137
+ try:
1138
+ infos = socket.getaddrinfo(host, None)
1139
+ except Exception as e: # noqa: BLE001
1140
+ raise ValueError(f"Cannot resolve host {host}: {e}")
1141
+ for info in infos:
1142
+ ip = info[4][0]
1143
+ try:
1144
+ addr = ipaddress.ip_address(ip)
1145
+ except ValueError:
1146
+ continue
1147
+ if (addr.is_private or addr.is_loopback or addr.is_link_local
1148
+ or addr.is_reserved or addr.is_multicast):
1149
+ raise ValueError(f"Refusing to fetch a private/internal address ({ip}).")
1150
+
1151
+
1152
+ def _rag_workdir(source_key: str) -> str:
1153
+ import os
1154
+ base = os.path.join(os.path.expanduser("~"), ".gonext", "rag-work", source_key)
1155
+ os.makedirs(base, exist_ok=True)
1156
+ return base
1157
+
1158
+
1159
+ def _rag_download(url: str, dest_path: str = "") -> str:
1160
+ """Download a public URL to a local path with SSRF + size guards. Returns the path."""
1161
+ import os
1162
+ src = _rag_gdrive_direct((url or "").strip())
1163
+ _rag_assert_safe_url(src)
1164
+ key = _rag_source_key(url)
1165
+ if not dest_path:
1166
+ name = os.path.basename(src.split("?")[0]) or "download.bin"
1167
+ if not name.lower().endswith(".zip") and "zip" in src.lower():
1168
+ name += ".zip"
1169
+ dest_path = os.path.join(_rag_workdir(key), name)
1170
+ os.makedirs(os.path.dirname(dest_path) or ".", exist_ok=True)
1171
+ req = urllib.request.Request(src, headers={"User-Agent": "gonext-rag/1.0"})
1172
+ total = 0
1173
+ with urllib.request.urlopen(req, timeout=120, context=_ssl_context()) as resp, \
1174
+ open(dest_path, "wb") as out:
1175
+ while True:
1176
+ block = resp.read(1024 * 256)
1177
+ if not block:
1178
+ break
1179
+ total += len(block)
1180
+ if total > _RAG_MAX_DOWNLOAD_BYTES:
1181
+ out.close()
1182
+ os.remove(dest_path)
1183
+ raise ValueError(
1184
+ f"Download exceeds the {_RAG_MAX_DOWNLOAD_BYTES // (1024*1024)} MB cap."
1185
+ )
1186
+ out.write(block)
1187
+ return dest_path
1188
+
1189
+
1190
+ def _rag_unzip(zip_path: str, dest_dir: str = "") -> tuple:
1191
+ """Extract a local zip with Zip-Slip + zip-bomb guards. Returns (dir, file_list)."""
1192
+ import os
1193
+ import zipfile
1194
+ if not dest_dir:
1195
+ dest_dir = os.path.splitext(zip_path)[0] + "_unzipped"
1196
+ os.makedirs(dest_dir, exist_ok=True)
1197
+ dest_root = os.path.realpath(dest_dir)
1198
+ written = 0
1199
+ names = []
1200
+ with zipfile.ZipFile(zip_path) as zf:
1201
+ infos = zf.infolist()
1202
+ if len(infos) > _RAG_MAX_FILES:
1203
+ raise ValueError(f"Zip has too many entries (> {_RAG_MAX_FILES}).")
1204
+ total_uncompressed = sum(i.file_size for i in infos)
1205
+ if total_uncompressed > _RAG_MAX_UNZIP_BYTES:
1206
+ raise ValueError("Zip uncompressed size exceeds the safety cap (possible zip bomb).")
1207
+ for info in infos:
1208
+ if info.is_dir():
1209
+ continue
1210
+ target = os.path.realpath(os.path.join(dest_dir, info.filename))
1211
+ # Zip Slip: extracted path must stay within dest_dir.
1212
+ if not (target == dest_root or target.startswith(dest_root + os.sep)):
1213
+ continue
1214
+ os.makedirs(os.path.dirname(target), exist_ok=True)
1215
+ with zf.open(info) as src, open(target, "wb") as dst:
1216
+ dst.write(src.read())
1217
+ written += 1
1218
+ names.append(info.filename)
1219
+ return dest_dir, names
1220
+
1221
+
1222
+ def _rag_iter_text_files(root: str):
1223
+ """Yield (abs_path, rel_path, ext) for indexable text files under root."""
1224
+ import os
1225
+ if os.path.isfile(root):
1226
+ yield root, os.path.basename(root), os.path.splitext(root)[1].lower()
1227
+ return
1228
+ for dirpath, dirnames, filenames in os.walk(root):
1229
+ dirnames[:] = [d for d in dirnames if d not in _RAG_SKIP_DIRS]
1230
+ for fn in filenames:
1231
+ ext = os.path.splitext(fn)[1].lower()
1232
+ base = fn.lower()
1233
+ if ext not in _RAG_TEXT_EXTS and base not in ("dockerfile", "makefile", ".gitignore"):
1234
+ continue
1235
+ abs_path = os.path.join(dirpath, fn)
1236
+ try:
1237
+ if os.path.getsize(abs_path) > _RAG_MAX_FILE_BYTES:
1238
+ continue
1239
+ except OSError:
1240
+ continue
1241
+ yield abs_path, os.path.relpath(abs_path, root), ext
1242
+
1243
+
1244
+ def _rag_chunk_text(text: str, rel_path: str, ext: str):
1245
+ """Line-aware windowed chunks with overlap. Returns list of {file,ext,start,end,text}."""
1246
+ if not text.strip():
1247
+ return []
1248
+ chunks = []
1249
+ n = len(text)
1250
+ step = max(1, _RAG_CHUNK_CHARS - _RAG_CHUNK_OVERLAP)
1251
+ i = 0
1252
+ while i < n:
1253
+ end = min(n, i + _RAG_CHUNK_CHARS)
1254
+ # Prefer to cut on a newline near the window end.
1255
+ if end < n:
1256
+ nl = text.rfind("\n", i + step, end)
1257
+ if nl > i:
1258
+ end = nl + 1
1259
+ piece = text[i:end].strip()
1260
+ if piece:
1261
+ chunks.append({"file": rel_path, "ext": ext, "start": i, "end": end, "text": piece})
1262
+ if end >= n:
1263
+ break
1264
+ i = end - _RAG_CHUNK_OVERLAP if end - _RAG_CHUNK_OVERLAP > i else end
1265
+ return chunks
1266
+
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("/")
1271
+ if not root:
1272
+ raise RuntimeError("No Ollama endpoint available for embeddings.")
1273
+ ctx = _ssl_context()
1274
+ # Preferred: batch /api/embed (recent Ollama).
1275
+ try:
1276
+ 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
+ )
1282
+ with urllib.request.urlopen(req, timeout=300, context=ctx) as resp:
1283
+ data = json.loads(resp.read().decode("utf-8", "replace"))
1284
+ embs = data.get("embeddings")
1285
+ if isinstance(embs, list) and embs and isinstance(embs[0], list):
1286
+ return embs
1287
+ 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).
1294
+ 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",
1301
+ )
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
1306
+
1307
+
1308
+ def _cosine(a: list, b: list) -> float:
1309
+ import math
1310
+ if not a or not b or len(a) != len(b):
1311
+ return -1.0
1312
+ dot = sum(x * y for x, y in zip(a, b))
1313
+ na = math.sqrt(sum(x * x for x in a))
1314
+ nb = math.sqrt(sum(y * y for y in b))
1315
+ if na == 0 or nb == 0:
1316
+ return -1.0
1317
+ return dot / (na * nb)
1318
+
1319
+
1320
+ def _rag_parse_location(loc: str):
1321
+ """Parse an ARN / s3:// URI / bucket[/prefix] into (bucket, prefix)."""
1322
+ s = (loc or "").strip() or "arn:aws:s3:::gonext-rag"
1323
+ rest = s
1324
+ m = re.match(r"^arn:aws:s3:::(.+)$", s, re.I)
1325
+ if m:
1326
+ rest = m.group(1)
1327
+ else:
1328
+ m = re.match(r"^s3://(.+)$", s, re.I)
1329
+ if m:
1330
+ rest = m.group(1)
1331
+ rest = rest.strip("/")
1332
+ if "/" in rest:
1333
+ bucket, prefix = rest.split("/", 1)
1334
+ else:
1335
+ bucket, prefix = rest, ""
1336
+ return bucket, prefix.strip("/")
1337
+
1338
+
1339
+ def _rag_s3_client(region: str, akid: str, secret: str):
1340
+ import boto3 # noqa: PLC0415
1341
+ return boto3.client(
1342
+ "s3", region_name=region or "ap-southeast-1",
1343
+ aws_access_key_id=akid, aws_secret_access_key=secret,
1344
+ )
1345
+
1346
+
1083
1347
  def run_agent_chat(cfg):
1084
1348
  try:
1085
1349
  from smolagents import CodeAgent, OpenAIServerModel, tool
@@ -1238,6 +1502,66 @@ def run_agent_chat(cfg):
1238
1502
  # action must never be silently available (same gating as the PDF reader).
1239
1503
  _EMAIL_AVAILABLE = bool(email_enabled and email_api_url and email_body_template)
1240
1504
 
1505
+ # ---- RAG config (download/unzip/index/search over a user's ZIP-at-URL) ----
1506
+ rag_enabled = bool(cfg.get("ragEnabled"))
1507
+ rag_location = (cfg.get("ragS3Location") or "arn:aws:s3:::gonext-rag").strip()
1508
+ rag_region = (cfg.get("ragAwsRegion") or "ap-southeast-1").strip()
1509
+ rag_akid = (cfg.get("ragAwsAccessKeyId") or "").strip()
1510
+ rag_secret = (cfg.get("ragAwsSecretAccessKey") or "").strip()
1511
+ rag_embed_model = (cfg.get("ragEmbedModel") or "nomic-embed-text").strip()
1512
+ try:
1513
+ rag_top_k = int(cfg.get("ragTopK") or 6)
1514
+ except (TypeError, ValueError):
1515
+ rag_top_k = 6
1516
+ 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("/"))
1520
+ try:
1521
+ import boto3 as _boto3_probe # noqa: F401,PLC0415
1522
+ _boto3_ok = True
1523
+ except Exception: # noqa: BLE001
1524
+ _boto3_ok = False
1525
+ _RAG_AVAILABLE = bool(rag_enabled and rag_akid and rag_secret and rag_ollama_root and _boto3_ok)
1526
+ if rag_enabled and not _boto3_ok:
1527
+ _log("RAG requested but boto3 is not installed in the worker python — RAG tools disabled")
1528
+
1529
+ def _rag_s3_and_loc():
1530
+ client = _rag_s3_client(rag_region, rag_akid, rag_secret)
1531
+ bucket, prefix = _rag_parse_location(rag_location)
1532
+ return client, bucket, prefix
1533
+
1534
+ def _rag_prefix_for(source_url: str, prefix: str) -> str:
1535
+ key = _rag_source_key(source_url)
1536
+ base = (prefix + "/") if prefix else ""
1537
+ return f"{base}agent-rag/{key}"
1538
+
1539
+ def _rag_load_chunks(client, bucket: str, base_prefix: str) -> list:
1540
+ """Load all chunk records (across shards) for a source. Returns [] if none."""
1541
+ out = []
1542
+ token = None
1543
+ while True:
1544
+ kw = {"Bucket": bucket, "Prefix": f"{base_prefix}/chunks"}
1545
+ if token:
1546
+ kw["ContinuationToken"] = token
1547
+ resp = client.list_objects_v2(**kw)
1548
+ for obj in resp.get("Contents", []) or []:
1549
+ if not obj["Key"].endswith(".jsonl"):
1550
+ continue
1551
+ body = client.get_object(Bucket=bucket, Key=obj["Key"])["Body"].read()
1552
+ for line in body.decode("utf-8", "replace").splitlines():
1553
+ line = line.strip()
1554
+ if line:
1555
+ try:
1556
+ out.append(json.loads(line))
1557
+ except json.JSONDecodeError:
1558
+ pass
1559
+ if resp.get("IsTruncated"):
1560
+ token = resp.get("NextContinuationToken")
1561
+ else:
1562
+ break
1563
+ return out
1564
+
1241
1565
  def _prior_email_preview() -> bool:
1242
1566
  """True if an earlier assistant turn showed a send_email preview awaiting confirm."""
1243
1567
  for m in reversed(messages[:last_user_idx]):
@@ -1384,7 +1708,25 @@ def run_agent_chat(cfg):
1384
1708
  if _EMAIL_AVAILABLE:
1385
1709
  _tool_num += 1
1386
1710
  _email_num = _tool_num
1711
+ if _RAG_AVAILABLE:
1712
+ _tool_num += 5 # download_file, unzip_file, rag_index, rag_add, rag_search
1387
1713
  _tool_count_word = {6: "SIX", 7: "SEVEN", 8: "EIGHT"}.get(_tool_num, str(_tool_num))
1714
+ _rag_tool_block = (
1715
+ " KNOWLEDGE BASE (RAG) — to SUMMARIZE or ANSWER questions about a ZIP of files at a URL:\n"
1716
+ " - download_file(url) — download a file (e.g. a .zip) to this machine.\n"
1717
+ " - 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"
1720
+ " - rag_add(source_url, text) — add extra info the user provides to that knowledge base.\n"
1721
+ " - rag_search(source_url, query) — retrieve the most relevant chunks to answer a question.\n"
1722
+ ) if _RAG_AVAILABLE else ""
1723
+ _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). "
1728
+ "If the user adds information, rag_add(source_url=url, text=...).\n"
1729
+ ) if _RAG_AVAILABLE else ""
1388
1730
  _pdf_read_tool_line = (
1389
1731
  f" {_pdf_num}. extract_text_from_pdf(url) — READ/summarize an EXISTING PDF file "
1390
1732
  "at a URL. This does NOT create a PDF (use create_pdf for that).\n"
@@ -1421,7 +1763,7 @@ def run_agent_chat(cfg):
1421
1763
  "Use this ONLY when the task explicitly asks for the date or time.\n"
1422
1764
  " 6. create_pdf(text, title='') — make a PDF document from text/data and return a "
1423
1765
  "download link. Use ONLY when the task asks to create/make/generate/export a PDF.\n"
1424
- + _pdf_read_tool_line + _email_tool_line +
1766
+ + _pdf_read_tool_line + _email_tool_line + _rag_tool_block +
1425
1767
  "\n"
1426
1768
  "http_request RETURN FORMAT: 'HTTP 200\\n{body}' — first line is 'HTTP <code>', body follows.\n"
1427
1769
  "\n"
@@ -1445,7 +1787,7 @@ def run_agent_chat(cfg):
1445
1787
  "- a PDF of info you must LOOK UP first (e.g. 'get the world cup schedule then make a "
1446
1788
  "pdf') -> FIRST web_search/fetch_url to gather the real facts, THEN in a later step "
1447
1789
  "create_pdf(text=<the gathered facts>, title=...). Never PDF the request wording itself.\n"
1448
- + _pdf_read_choose_line + _email_choose_line +
1790
+ + _pdf_read_choose_line + _email_choose_line + _rag_choose_line +
1449
1791
  "- A specific known API/URL was given -> http_request().\n"
1450
1792
  "\n"
1451
1793
  "RULES:\n"
@@ -2275,14 +2617,160 @@ def run_agent_chat(cfg):
2275
2617
  client_kwargs={"timeout": 600.0, "max_retries": 0},
2276
2618
  **extra_model_kwargs,
2277
2619
  )
2620
+ @tool
2621
+ def download_file(url: str, dest_path: str = "") -> str:
2622
+ """Download a file (e.g. a .zip) from a public http(s) URL to a local path on this machine. Use this FIRST when the user gives a URL to a zip of files; then unzip_file, then rag_index.
2623
+
2624
+ Args:
2625
+ url: public http(s) URL of the file to download.
2626
+ dest_path: optional local save path; leave empty for an automatic work-dir path.
2627
+ """
2628
+ _emit({"type": "step", "text": f"Downloading → {url[:80]}"})
2629
+ try:
2630
+ import os
2631
+ path = _rag_download(url, (dest_path or "").strip())
2632
+ return f"Downloaded to {path} ({os.path.getsize(path)} bytes). Next: unzip_file(zip_path='{path}') if it is a zip."
2633
+ except Exception as e: # noqa: BLE001
2634
+ return f"Error: download failed: {type(e).__name__}: {e}"
2635
+
2636
+ @tool
2637
+ def unzip_file(zip_path: str, dest_dir: str = "") -> str:
2638
+ """Extract a local .zip file into a local directory. Use after download_file. Returns the directory and a short file listing.
2639
+
2640
+ Args:
2641
+ zip_path: local path to the .zip (from download_file).
2642
+ dest_dir: optional target directory; leave empty for an automatic one.
2643
+ """
2644
+ _emit({"type": "step", "text": "Unzipping…"})
2645
+ try:
2646
+ d, names = _rag_unzip((zip_path or "").strip(), (dest_dir or "").strip())
2647
+ preview = ", ".join(names[:15])
2648
+ more = f" …(+{len(names) - 15} more)" if len(names) > 15 else ""
2649
+ return f"Unzipped {len(names)} files to {d}. Files: {preview}{more}. Next: rag_index(path='{d}', source_url=<the original url>)."
2650
+ except Exception as e: # noqa: BLE001
2651
+ return f"Error: unzip failed: {type(e).__name__}: {e}"
2652
+
2653
+ @tool
2654
+ def rag_index(path: str, source_url: str) -> str:
2655
+ """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.
2656
+
2657
+ Args:
2658
+ path: local file or directory to index (from unzip_file).
2659
+ source_url: the original URL the user gave — used as the knowledge-base key.
2660
+ """
2661
+ if not _RAG_AVAILABLE:
2662
+ return "Error: RAG is not configured (enable it + add AWS credentials in Settings)."
2663
+ _emit({"type": "step", "text": f"RAG source → {source_url}"})
2664
+ _emit({"type": "step", "text": "Indexing files for RAG…"})
2665
+ try:
2666
+ import time as _t
2667
+ records = []
2668
+ files = 0
2669
+ for abs_path, rel, ext in _rag_iter_text_files((path or "").strip()):
2670
+ try:
2671
+ with open(abs_path, "r", encoding="utf-8", errors="replace") as fh:
2672
+ text = fh.read()
2673
+ except OSError:
2674
+ continue
2675
+ files += 1
2676
+ records.extend(_rag_chunk_text(text, rel, ext))
2677
+ if not records:
2678
+ return "No indexable text files found at that path."
2679
+ for i in range(0, len(records), 64):
2680
+ batch = records[i:i + 64]
2681
+ vecs = _ollama_embed(rag_ollama_root, rag_embed_model, [r["text"] for r in batch])
2682
+ for r, v in zip(batch, vecs):
2683
+ r["embedding"] = v
2684
+ _emit({"type": "step", "text": f"Embedded {min(i + 64, len(records))}/{len(records)} chunks…"})
2685
+ client, bucket, prefix = _rag_s3_and_loc()
2686
+ base = _rag_prefix_for(source_url, prefix)
2687
+ shard = f"{base}/chunks-{int(_t.time())}.jsonl"
2688
+ body = "\n".join(json.dumps(r, ensure_ascii=False) for r in records).encode("utf-8")
2689
+ client.put_object(Bucket=bucket, Key=shard, Body=body, ContentType="application/x-ndjson")
2690
+ manifest = {
2691
+ "sourceUrl": source_url, "embedModel": rag_embed_model,
2692
+ "dims": len(records[0].get("embedding") or []),
2693
+ "files": files, "chunks": len(records),
2694
+ "updatedAt": _t.strftime("%Y-%m-%dT%H:%M:%SZ", _t.gmtime()),
2695
+ }
2696
+ client.put_object(Bucket=bucket, Key=f"{base}/manifest.json",
2697
+ Body=json.dumps(manifest).encode("utf-8"),
2698
+ ContentType="application/json")
2699
+ return (f"Indexed {files} files into {len(records)} chunks for {source_url}. "
2700
+ f"Use rag_search(source_url='{source_url}', query=...) to retrieve, then answer the user.")
2701
+ except Exception as e: # noqa: BLE001
2702
+ return f"Error: indexing failed: {type(e).__name__}: {e}"
2703
+
2704
+ @tool
2705
+ def rag_add(source_url: str, text: str) -> str:
2706
+ """Add extra text/notes to an EXISTING knowledge base (keyed by source_url). Use when the user gives more information to remember for later questions.
2707
+
2708
+ Args:
2709
+ source_url: the knowledge-base key (the original URL).
2710
+ text: the new information to embed and store.
2711
+ """
2712
+ if not _RAG_AVAILABLE:
2713
+ return "Error: RAG is not configured."
2714
+ if not (text or "").strip():
2715
+ return "Error: no text to add."
2716
+ _emit({"type": "step", "text": f"RAG source → {source_url}"})
2717
+ _emit({"type": "step", "text": "Adding info to RAG…"})
2718
+ try:
2719
+ import time as _t
2720
+ records = _rag_chunk_text(text, "user-note", ".txt")
2721
+ vecs = _ollama_embed(rag_ollama_root, rag_embed_model, [r["text"] for r in records])
2722
+ for r, v in zip(records, vecs):
2723
+ r["embedding"] = v
2724
+ client, bucket, prefix = _rag_s3_and_loc()
2725
+ base = _rag_prefix_for(source_url, prefix)
2726
+ shard = f"{base}/chunks-add-{int(_t.time())}.jsonl"
2727
+ body = "\n".join(json.dumps(r, ensure_ascii=False) for r in records).encode("utf-8")
2728
+ client.put_object(Bucket=bucket, Key=shard, Body=body, ContentType="application/x-ndjson")
2729
+ return f"Added {len(records)} chunk(s) to the knowledge base for {source_url}."
2730
+ except Exception as e: # noqa: BLE001
2731
+ return f"Error: rag_add failed: {type(e).__name__}: {e}"
2732
+
2733
+ @tool
2734
+ def rag_search(source_url: str, query: str, k: int = 0) -> str:
2735
+ """Retrieve the most relevant chunks from the knowledge base for source_url. Use to ANSWER questions about the indexed files; write your answer from the returned chunks.
2736
+
2737
+ Args:
2738
+ source_url: the knowledge-base key (the original URL).
2739
+ query: what to look for.
2740
+ k: number of chunks to return (0 = configured default).
2741
+ """
2742
+ if not _RAG_AVAILABLE:
2743
+ return "Error: RAG is not configured."
2744
+ _emit({"type": "step", "text": f"RAG source → {source_url}"})
2745
+ _emit({"type": "step", "text": f"Searching RAG → {query[:60]}"})
2746
+ try:
2747
+ client, bucket, prefix = _rag_s3_and_loc()
2748
+ base = _rag_prefix_for(source_url, prefix)
2749
+ chunks = _rag_load_chunks(client, bucket, base)
2750
+ if not chunks:
2751
+ 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]
2753
+ scored = sorted(
2754
+ ((_cosine(qvec, c.get("embedding") or []), c) for c in chunks),
2755
+ key=lambda x: x[0], reverse=True,
2756
+ )
2757
+ topk = scored[: (k if k and k > 0 else rag_top_k)]
2758
+ parts = [f"[{c.get('file', '?')}] (score {score:.2f})\n{c.get('text', '')}"
2759
+ for score, c in topk]
2760
+ return "\n\n---\n\n".join(parts) if parts else "No matches."
2761
+ except Exception as e: # noqa: BLE001
2762
+ return f"Error: rag_search failed: {type(e).__name__}: {e}"
2763
+
2278
2764
  agent_tools = [http_request, web_search, fetch_url, calculate,
2279
2765
  get_current_datetime, create_pdf]
2280
- # Only register the PDF reader / email tools when available (kept in sync with
2281
- # the gated hint entries above so the model never sees a tool it can't run).
2766
+ # Only register the PDF reader / email / RAG tools when available (kept in sync
2767
+ # with the gated hint entries above so the model never sees a tool it can't run).
2282
2768
  if _PDF_READ_AVAILABLE:
2283
2769
  agent_tools.append(extract_text_from_pdf)
2284
2770
  if _EMAIL_AVAILABLE:
2285
2771
  agent_tools.append(send_email)
2772
+ if _RAG_AVAILABLE:
2773
+ agent_tools += [download_file, unzip_file, rag_index, rag_add, rag_search]
2286
2774
  agent_kwargs = dict(
2287
2775
  tools=agent_tools,
2288
2776
  model=model,
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@tiens.nguyen/gonext-local-worker",
3
- "version": "1.0.153",
3
+ "version": "1.0.157",
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",