"""Asynchronous, file-backed assistant export and ordinary Wiki restore.""" from __future__ import annotations import asyncio from pathlib import Path from uuid import uuid4 from fastapi import APIRouter, File, HTTPException, Request, UploadFile from fastapi.responses import FileResponse from app.gateway.llmwiki_index_scheduler import spawn_local_wiki_index_job from app.gateway.routers.assistant_knowledge import _actor, _list_export_manifests, _package_root, _read_export_manifest, _store, _write_export_manifest from deerflow.assistant_knowledge.archive import PackageArchive from deerflow.assistant_knowledge.package_io import package_path, write_jsonl_record from deerflow.integrations.weknora.local_index.normalize import normalize_wiki_page from deerflow.integrations.weknora.runtime import build_weknora_client, get_resolved_llmwiki_runtime from deerflow.persistence.assistant_knowledge.store import _vector_blob_from_payload router = APIRouter(tags=["wiki-packages"]) def job_view(row): if row["mapping_id"].startswith("assistant:"): base_id = row["mapping_id"].split(":", 1)[1] row["package_download_url"] = f"/api/assistant-knowledge/bases/{base_id}/export-package/{row['id']}/download" elif row["mapping_id"].startswith("restore:"): row["package_download_url"] = None return row async def export_assistant(app, base_id, job_id): scope = f"assistant:{base_id}" path = package_path(_package_root(), job_id) temp = path.with_suffix(".tmp") counts = {} try: _write_export_manifest(job_id=job_id, mapping_id=scope, status="running", phase="exporting") with temp.open("w", encoding="utf-8", newline="\n") as handle: async for kind, payload in app.state.assistant_knowledge_store.iter_export_records(base_id): await asyncio.to_thread(write_jsonl_record, handle, kind, payload) counts[kind] = counts.get(kind, 0) + 1 if counts.get("wiki_page") and not counts.get("vector"): raise ValueError("该助手知识库尚无可导出的 Wiki 向量,请先从已向量化的普通知识库同步") await asyncio.to_thread(temp.replace, path) _write_export_manifest(job_id=job_id, mapping_id=scope, status="completed", phase="package_ready", counts=counts) except Exception as exc: _write_export_manifest(job_id=job_id, mapping_id=scope, status="failed", phase="failed", counts=counts, error=str(exc)) @router.post("/api/assistant-knowledge/bases/{base_id}/export-package", status_code=202) async def create_export(request: Request, base_id: str): await _actor(request, admin=True) if not await _store(request).get_base(base_id): raise HTTPException(404, "助手知识库不存在") job_id = str(uuid4()) row = _write_export_manifest(job_id=job_id, mapping_id=f"assistant:{base_id}", status="queued", phase="queued") spawn_local_wiki_index_job(request.app, export_assistant(request.app, base_id, job_id)) return job_view(row) @router.get("/api/assistant-knowledge/bases/{base_id}/export-package") async def list_exports(request: Request, base_id: str): await _actor(request, admin=True) return {"jobs": [job_view(row) for row in await asyncio.to_thread(_list_export_manifests, f"assistant:{base_id}", limit=100)]} @router.get("/api/assistant-knowledge/bases/{base_id}/export-package/{job_id}/download") async def download_export(request: Request, base_id: str, job_id: str): await _actor(request, admin=True) row = _read_export_manifest(job_id) if not row or row["mapping_id"] != f"assistant:{base_id}" or row["status"] != "completed" or not row["file_exists"]: raise HTTPException(404, "导出包尚未就绪") return FileResponse(package_path(_package_root(), job_id), media_type="application/x-ndjson", filename=f"wiki-vector-{job_id}.jsonl") async def restore_package(app, mapping, job_id, path: Path): scope = f"restore:{mapping['id']}" archive = None counts = {"wiki_page": 0, "vector": 0} store = app.state.llmwiki_index_store owner = f"package:{job_id}" leased = False try: _write_export_manifest(job_id=job_id, mapping_id=scope, status="running", phase="validating") archive = await asyncio.to_thread(PackageArchive, path) embedding = app.state.llmwiki_embedding if embedding is None: raise ValueError("请配置与数据包相同的 Wiki 编码模型,以便复用向量进行检索") if embedding.dimensions is None: await embedding.embed_query("Wiki embedding dimension probe") # Validate the complete package before the first remote mutation. slugs = {p["slug"] for p in archive["wiki_pages"]} vector_slugs = set() for raw in archive["vectors"]: slug = str(raw.get("wiki_slug") or raw.get("page_slug") or raw.get("slug") or "") if slug not in slugs or _vector_blob_from_payload(raw) is None: raise ValueError(f"向量数据或 Wiki 引用无效:{slug}") if raw.get("embedding_fingerprint") != embedding.fingerprint or raw.get("embedding_dimensions") != embedding.dimensions: raise ValueError("数据包编码模型指纹与当前配置不一致;请使用同一模型配置,禁止静默重新编码") vector_slugs.add(slug) if not vector_slugs: raise ValueError("数据包中没有 Wiki 向量") for page in archive["wiki_pages"]: if page["slug"] not in vector_slugs and page.get("page_type") not in {"index", "directory"} and (page.get("content") or "").strip(): raise ValueError(f"Wiki 缺少向量:{page['slug']}") leased = await store.try_acquire_sync_lease(mapping["id"], owner=owner, sync_id=job_id, lease_seconds=900, fingerprint=embedding.fingerprint) if not leased: raise ValueError("目标知识库正在同步,请稍后重试") client = build_weknora_client(get_resolved_llmwiki_runtime(app.state.config)) _write_export_manifest(job_id=job_id, mapping_id=scope, status="running", phase="restoring") folders = {(): ""} for raw in archive["wiki_pages"]: if not await store.renew_sync_lease(mapping["id"], owner=owner, lease_seconds=900): raise RuntimeError("导入租约已失效,请重新发起") # Restore the provider's Wiki schema, never raw documents; raw # uploads would trigger generation and embedding a second time. categories = raw.get("category_path") or [] if isinstance(categories, str): categories = [part for part in categories.split("/") if part] parent = () for name in categories: path_key = (*parent, name) if path_key not in folders: folder = await client.create_wiki_folder(mapping["weknora_id"], name=name, parent_id=folders[parent]) folders[path_key] = folder["id"] parent = path_key payload = {key: value for key, value in raw.items() if key not in {"id", "tenant_id", "knowledge_base_id", "folder_id", "origins"}} payload["folder_id"] = folders[parent] if raw.get("tags"): payload["page_metadata"] = {**(raw.get("page_metadata") or {}), "tags": raw["tags"]} remote = await client.create_wiki_page(mapping["weknora_id"], payload) page = normalize_wiki_page(mapping["id"], {**raw, **remote}) vectors = [] for value in archive.vectors_for(raw["slug"]): blob, dimensions = _vector_blob_from_payload(value) vectors.append({**value, "vector_blob": blob, "embedding_dimensions": dimensions}) await store.replace_page_vectors(page=page, vectors=vectors, fingerprint=embedding.fingerprint, sync_id=job_id) counts["wiki_page"] += 1 counts["vector"] += len(vectors) _write_export_manifest(job_id=job_id, mapping_id=scope, status="running", phase="restoring", counts=counts) await client.rebuild_wiki_links(mapping["weknora_id"]) await store.finish_sync(mapping["id"], owner=owner, sync_id=job_id, success=True, remote_page_count=counts["wiki_page"]) leased = False if app.state.llmwiki_vector_cache: app.state.llmwiki_vector_cache.invalidate(mapping["id"]) from app.gateway.knowledge_transfer import sync_to_global await sync_to_global(app, mapping) _write_export_manifest(job_id=job_id, mapping_id=scope, status="completed", phase="completed", counts=counts) except Exception as exc: if leased: await store.finish_sync(mapping["id"], owner=owner, sync_id=job_id, success=False, remote_page_count=counts["wiki_page"], error=str(exc)) _write_export_manifest(job_id=job_id, mapping_id=scope, status="failed", phase="failed", counts=counts, error=str(exc)) finally: if archive is not None: await asyncio.to_thread(archive.close) @router.post("/api/llmwiki/knowledge-bases/{mapping_id}/import-package", status_code=202) async def upload_restore(request: Request, mapping_id: str, file: UploadFile = File(...)): actor = await _actor(request, admin=True) mapping = await request.app.state.llmwiki_store.get_authorized(mapping_id, actor or "system", write=True, is_admin=True) if not mapping or mapping.get("is_conversation_deposit") or mapping.get("owner_user_id") == "system": raise HTTPException(404, "目标普通知识库不存在或不可导入") client = build_weknora_client(get_resolved_llmwiki_runtime(request.app.state.config)) remote = await client.get_knowledge_base(mapping["weknora_id"]) if not (remote.get("indexing_strategy") or {}).get("wiki_enabled"): raise HTTPException(422, "请先开启目标知识库的 Wiki") existing = await client.list_wiki_pages(mapping["weknora_id"], page_size=1) if existing.get("total"): raise HTTPException(409, "请导入到新建的空 Wiki 知识库,避免覆盖已有页面") job_id = str(uuid4()) scope = f"restore:{mapping_id}" row = _write_export_manifest(job_id=job_id, mapping_id=scope, status="queued", phase="uploading") path = package_path(_package_root(), job_id) try: with path.open("wb") as handle: while data := await file.read(1024 * 1024): await asyncio.to_thread(handle.write, data) except Exception as exc: _write_export_manifest(job_id=job_id, mapping_id=scope, status="failed", phase="upload_failed", error=str(exc)) raise finally: await file.close() spawn_local_wiki_index_job(request.app, restore_package(request.app, mapping, job_id, path)) return job_view(row) @router.get("/api/llmwiki/knowledge-bases/{mapping_id}/import-package") async def list_restores(request: Request, mapping_id: str): await _actor(request, admin=True) return {"jobs": [job_view(row) for row in await asyncio.to_thread(_list_export_manifests, f"restore:{mapping_id}", limit=100)]}