118 lines
3.8 KiB
Python
118 lines
3.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Walk the intranet/ folder and ingest every text-bearing file into AnythingLLM.
|
|
|
|
Pipeline per file:
|
|
1. POST /api/v1/document/upload -> AnythingLLM parses + chunks + saves
|
|
metadata, returns documents:[{location}]
|
|
2. Collect all returned locations.
|
|
3. POST /api/v1/workspace/intranet/update-embeddings with adds=locations
|
|
-> AnythingLLM embeds each chunk via Ollama bge-m3 and indexes into LanceDB.
|
|
|
|
We do step 1 sequentially (server processes serially anyway), then a single
|
|
bulk step 2 so the heavy embedding pass is one operation we can watch.
|
|
|
|
Skipped extensions: jpg/png (no OCR), mp4 (no whisper for now), crt, vsdx.
|
|
"""
|
|
import os
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import requests
|
|
|
|
BASE = os.environ.get("ANYTHINGLLM_URL", "http://127.0.0.1:3035")
|
|
API_KEY = os.environ.get("API_KEY",
|
|
"AILM-78c5154c6dc4fc04a940784c8ad9927e8882c7e16bf4cdcd044ccd0fd00af168")
|
|
WORKSPACE = os.environ.get("WORKSPACE", "intranet")
|
|
ROOT = Path(os.environ.get("INTRANET_ROOT",
|
|
"/home/klas/Prace/AI/portal/anythingllm/intranet/INTRANET"))
|
|
|
|
ALLOWED_EXTS = {".pdf", ".docx", ".doc", ".xlsx", ".xls",
|
|
".pptx", ".txt", ".html", ".md"}
|
|
|
|
HEADERS = {"Authorization": f"Bearer {API_KEY}"}
|
|
|
|
|
|
def upload_file(path: Path) -> str | None:
|
|
"""Upload one file. Returns its document 'location' on success."""
|
|
with open(path, "rb") as f:
|
|
files = {"file": (path.name, f, "application/octet-stream")}
|
|
r = requests.post(
|
|
f"{BASE}/api/v1/document/upload",
|
|
headers=HEADERS,
|
|
files=files,
|
|
data={"addToWorkspaces": ""}, # we'll attach in bulk after
|
|
timeout=300,
|
|
)
|
|
if not r.ok:
|
|
print(f" ! HTTP {r.status_code}: {r.text[:200]}")
|
|
return None
|
|
data = r.json()
|
|
if not data.get("success"):
|
|
print(f" ! upload failed: {data.get('error', data)}")
|
|
return None
|
|
docs = data.get("documents") or []
|
|
if not docs:
|
|
print(" ! no documents in response")
|
|
return None
|
|
return docs[0]["location"]
|
|
|
|
|
|
def attach_to_workspace(locations: list[str]) -> dict:
|
|
"""Embed + index everything we uploaded into the workspace."""
|
|
r = requests.post(
|
|
f"{BASE}/api/v1/workspace/{WORKSPACE}/update-embeddings",
|
|
headers={**HEADERS, "Content-Type": "application/json"},
|
|
json={"adds": locations, "deletes": []},
|
|
timeout=3600,
|
|
)
|
|
r.raise_for_status()
|
|
return r.json()
|
|
|
|
|
|
def main() -> int:
|
|
if not ROOT.exists():
|
|
print(f"Root does not exist: {ROOT}", file=sys.stderr)
|
|
return 1
|
|
|
|
files = sorted(
|
|
p for p in ROOT.rglob("*")
|
|
if p.is_file() and p.suffix.lower() in ALLOWED_EXTS
|
|
)
|
|
print(f"Found {len(files)} ingestible files under {ROOT}")
|
|
|
|
locations: list[str] = []
|
|
failed: list[Path] = []
|
|
t0 = time.time()
|
|
for i, path in enumerate(files, 1):
|
|
rel = path.relative_to(ROOT)
|
|
size_kb = path.stat().st_size / 1024
|
|
print(f"[{i:>3}/{len(files)}] {rel} ({size_kb:.0f} KB)")
|
|
loc = upload_file(path)
|
|
if loc:
|
|
locations.append(loc)
|
|
else:
|
|
failed.append(path)
|
|
|
|
print(f"\nUploaded {len(locations)}/{len(files)} files in "
|
|
f"{time.time()-t0:.0f}s. Failed: {len(failed)}")
|
|
if failed:
|
|
print("Failed files:")
|
|
for p in failed:
|
|
print(f" - {p.relative_to(ROOT)}")
|
|
|
|
if not locations:
|
|
print("Nothing to embed.")
|
|
return 1
|
|
|
|
print(f"\nEmbedding {len(locations)} docs into workspace '{WORKSPACE}'...")
|
|
t0 = time.time()
|
|
result = attach_to_workspace(locations)
|
|
print(f"Done in {time.time()-t0:.0f}s")
|
|
print(f"Workspace now has {len(result.get('workspace', {}).get('documents', []))} documents")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|