153 lines
5.5 KiB
Python
153 lines
5.5 KiB
Python
"""File-backed assistant knowledge import/export packages.
|
|
|
|
The assistant knowledge base is fed by WeKnora-exported data. Keep that
|
|
boundary explicit: a WeKnora export is first materialized as a durable JSONL
|
|
package, then imported asynchronously from that package. JSONL keeps large
|
|
exports append-friendly and avoids treating the HTTP request body as the job
|
|
state.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from collections.abc import Iterable, Mapping
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
PACKAGE_VERSION = 1
|
|
PACKAGE_RECORD_TYPES = {
|
|
"source",
|
|
"wiki_page",
|
|
"document",
|
|
"chunk",
|
|
"vector",
|
|
"entity",
|
|
"relation",
|
|
"graph",
|
|
}
|
|
|
|
_PACKAGE_SNIFF_BYTES = 256 * 1024
|
|
|
|
|
|
def looks_like_knowledge_package(content: bytes) -> bool:
|
|
"""Recognize a DeerFlow Wiki-vector JSONL package from a bounded prefix.
|
|
|
|
The same check is used at ordinary-document upload boundaries so an
|
|
exported package can never be handed to a document parser (and therefore
|
|
to a model) by mistake. Source metadata may be split across consecutive
|
|
``source`` records for compatibility with packages produced by older test
|
|
helpers and releases.
|
|
"""
|
|
|
|
source: dict[str, Any] = {}
|
|
text = content[:_PACKAGE_SNIFF_BYTES].decode("utf-8-sig", errors="ignore")
|
|
for raw_line in text.splitlines():
|
|
# In a multipart upload the file's BOM is not the request body's BOM,
|
|
# so strip it again at the beginning of each candidate line.
|
|
line = raw_line.strip().lstrip("\ufeff").strip()
|
|
if not line.startswith("{"):
|
|
continue
|
|
try:
|
|
record = json.loads(line)
|
|
except (TypeError, ValueError):
|
|
continue
|
|
if not isinstance(record, dict):
|
|
continue
|
|
if record.get("record_type") != "source":
|
|
# A real package starts with one or more source records. Do not
|
|
# scan arbitrary JSONL prose for a package-shaped example later.
|
|
return False
|
|
payload = record.get("payload")
|
|
if not isinstance(payload, dict):
|
|
continue
|
|
source.update(payload)
|
|
provider = source.get("provider") or source.get("type")
|
|
if (
|
|
source.get("package_version") == PACKAGE_VERSION
|
|
and provider in {"weknora", "assistant"}
|
|
):
|
|
return True
|
|
return False
|
|
|
|
|
|
def package_path(root: Path, job_id: str) -> Path:
|
|
safe = "".join(ch for ch in str(job_id) if ch.isalnum() or ch in {"-", "_"})
|
|
return root / f"{safe}.jsonl"
|
|
|
|
|
|
def write_jsonl_record(handle: Any, record_type: str, payload: Mapping[str, Any]) -> None:
|
|
handle.write(json.dumps({"record_type": record_type, "payload": dict(payload)}, ensure_ascii=False))
|
|
handle.write("\n")
|
|
|
|
|
|
def write_package(path: Path, records: Iterable[tuple[str, Mapping[str, Any]]]) -> dict[str, int]:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
counts: dict[str, int] = {key: 0 for key in PACKAGE_RECORD_TYPES}
|
|
tmp_path = path.with_suffix(path.suffix + ".tmp")
|
|
with tmp_path.open("w", encoding="utf-8", newline="\n") as handle:
|
|
write_jsonl_record(
|
|
handle,
|
|
"source",
|
|
{
|
|
"package_version": PACKAGE_VERSION,
|
|
},
|
|
)
|
|
counts["source"] += 1
|
|
for record_type, payload in records:
|
|
if record_type not in PACKAGE_RECORD_TYPES:
|
|
continue
|
|
write_jsonl_record(handle, record_type, payload)
|
|
counts[record_type] = counts.get(record_type, 0) + 1
|
|
tmp_path.replace(path)
|
|
return counts
|
|
|
|
|
|
def read_package(path: Path) -> dict[str, Any]:
|
|
package: dict[str, Any] = {
|
|
"source": {},
|
|
"wiki_pages": [],
|
|
"documents": [],
|
|
"chunks": [],
|
|
"vectors": [],
|
|
"entities": [],
|
|
"relations": [],
|
|
"graph": {"nodes": [], "edges": []},
|
|
}
|
|
with path.open("r", encoding="utf-8") as handle:
|
|
for line in handle:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
record = json.loads(line)
|
|
if not isinstance(record, dict):
|
|
continue
|
|
record_type = str(record.get("record_type") or "")
|
|
payload = record.get("payload")
|
|
if not isinstance(payload, dict):
|
|
continue
|
|
if record_type == "source":
|
|
package["source"] = {**package.get("source", {}), **payload}
|
|
elif record_type == "wiki_page":
|
|
package["wiki_pages"].append(payload)
|
|
elif record_type == "document":
|
|
package["documents"].append(payload)
|
|
elif record_type == "chunk":
|
|
package["chunks"].append(payload)
|
|
elif record_type == "vector":
|
|
package["vectors"].append(payload)
|
|
elif record_type == "entity":
|
|
package["entities"].append(payload)
|
|
elif record_type == "relation":
|
|
package["relations"].append(payload)
|
|
elif record_type == "graph":
|
|
nodes = payload.get("nodes")
|
|
edges = payload.get("edges")
|
|
if isinstance(nodes, list):
|
|
package["graph"]["nodes"].extend(item for item in nodes if isinstance(item, dict))
|
|
if isinstance(edges, list):
|
|
package["graph"]["edges"].extend(item for item in edges if isinstance(item, dict))
|
|
for key, value in payload.items():
|
|
if key not in {"nodes", "edges"}:
|
|
package["graph"][key] = value
|
|
return package
|