import base64 import json import sys from io import BytesIO from types import SimpleNamespace import pytest from fastapi import FastAPI, UploadFile from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from app.gateway.routers import wiki_packages from deerflow.assistant_knowledge.archive import PackageArchive from deerflow.assistant_knowledge.package_io import write_package from deerflow.config.llmwiki_config import LocalWikiEmbeddingConfig from deerflow.integrations.weknora.local_index import embedding as embedding_module from deerflow.integrations.weknora.local_index.chunking import chunk_wiki_page from deerflow.integrations.weknora.local_index.embedding import StrictWikiEmbeddingClient from deerflow.integrations.weknora.local_index.normalize import normalize_wiki_page from deerflow.integrations.weknora.local_index.vector_codec import encode_vector from deerflow.persistence.assistant_knowledge import AssistantKnowledgeRepository from deerflow.persistence.base import Base from deerflow.persistence.llmwiki_index.memory import MemoryLlmWikiIndexStore @pytest.mark.asyncio async def test_metadata_vectors_and_links_survive_roundtrip_without_encoding(tmp_path, monkeypatch): engine = create_async_engine("sqlite+aiosqlite:///:memory:") async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) store = AssistantKnowledgeRepository(async_sessionmaker(engine, expire_on_commit=False)) blob, _ = encode_vector([0.3, 0.5, 0.7]) pages = [ { "slug": "concept/cooling", "title": "散热", "page_type": "concept", "content": "散热管理用于降低设备温度。参见 [[entity/fan]]", "summary": "设备散热", "tags": ["设备", "维护"], "category_path": ["运维", "硬件"], "aliases": ["降温"], "out_links": ["entity/fan"], "in_links": [], }, {"slug": "entity/fan", "title": "风扇", "page_type": "entity", "content": "风扇通过气流帮助散热，参见 [[concept/cooling]]。", "aliases": [], "out_links": ["concept/cooling"], "in_links": ["concept/cooling"]}, ] vectors = [ { "wiki_slug": p["slug"], "section_index": 0, "section_content": p["content"], "content_hash": "test", "embedding_model": "test-model", "embedding_fingerprint": "test-fp", "embedding_dimensions": 3, "vector_blob_base64": base64.b64encode(blob).decode(), } for p in pages ] original = tmp_path / "original.jsonl" entities = [{"id": p["slug"], "name": p["title"], "type": p["page_type"], "source_wiki_slug": p["slug"]} for p in pages] relations = [{"source_entity_id": pages[0]["slug"], "target_entity_id": pages[1]["slug"], "source_wiki_slug": pages[0]["slug"], "predicate": "references"}] write_package(original, [ ("source", {"provider": "weknora", "package_version": 1}), *[("wiki_page", p) for p in pages], *[("vector", v) for v in vectors], *[("entity", e) for e in entities], *[("relation", r) for r in relations], ("chunk", {"id": "raw-chunk", "content": "原始文档分段不是 Wiki 页面"}), ]) archive = PackageArchive(original) base = await store.initialize(created_by="admin") job = await store.import_package(package=archive, source_type="weknora", source_key="source", source_name="来源", trigger="test", created_by="admin") assert job["counts"]["vectors"] == 2 assert job["counts"]["wiki_pages"] == 2 assert job["counts"]["chunk_pages"] == 1 imported = await store.get_page(pages[0]["slug"]) for field in ("category_path", "tags", "aliases", "out_links"): assert imported[field] == pages[0][field] assert imported["content_markdown"] == pages[0]["content"] # The global base may receive the same Wiki through multiple libraries. # Only the current article's import supplies vectors; older contributions # remain auditable but must not duplicate/stale the restored index. await store.import_package(package=archive, source_type="weknora", source_key="second-source", source_name="第二来源", trigger="test", created_by="admin") exported = tmp_path / "export.jsonl" write_package(exported, [record async for record in store.iter_export_records(base["id"])]) restored_archive = PackageArchive(exported) assert {v["vector_blob_base64"] for v in restored_archive["vectors"]} == {base64.b64encode(blob).decode()} assert sorted((v["wiki_slug"], v["section_index"]) for v in restored_archive["vectors"]) == sorted((v["wiki_slug"], v["section_index"]) for v in vectors) assert all(r["source_wiki_slug"] in {p["slug"] for p in pages} for r in restored_archive["relations"]) index = MemoryLlmWikiIndexStore() received = [] class Remote: async def create_wiki_folder(self, kb, *, name, parent_id=""): return {"id": f"{parent_id}/{name}"} async def create_wiki_page(self, kb, page): received.append(page) return page async def rebuild_wiki_links(self, kb): pass class Embedding: dimensions = 3 fingerprint = "test-fp" async def embed_query(self, query): raise AssertionError("Existing package must not be encoded") app = FastAPI() app.state.llmwiki_index_store = index app.state.llmwiki_embedding = Embedding() app.state.llmwiki_vector_cache = None app.state.config = None monkeypatch.setattr(wiki_packages, "get_resolved_llmwiki_runtime", lambda c: None) monkeypatch.setattr(wiki_packages, "build_weknora_client", lambda c: Remote()) monkeypatch.setattr(wiki_packages, "_package_root", lambda: tmp_path) from app.gateway import knowledge_transfer async def sync(app, mapping): return {"status": "completed"} monkeypatch.setattr(knowledge_transfer, "sync_to_global", sync) await wiki_packages.restore_package(app, {"id": "target", "weknora_id": "remote"}, "job-test", exported) for raw in received: assert raw["content"] == next(p["content"] for p in pages if p["slug"] == raw["slug"]) local = await index.get_page("target", raw["slug"]) assert local["content_hash"] == normalize_wiki_page("target", raw)["content_hash"] snapshot = await index.load_vector_snapshot("target", "test-fp") assert len(snapshot["rows"]) == 2 assert all(r["vector_blob"] == blob for r in snapshot["rows"]) class SearchEmbedding: dimensions = 3 fingerprint = "test-fp" async def embed_query(self, query): return [1, 0, 0] found = await store.search_pages(base_ids=[base["id"]], query="没有关键词重合的查询", embedding_client=SearchEmbedding()) assert found and all(r["retrieval_mode"] == "vector" for r in found) SearchEmbedding.fingerprint = "different-model-same-dimension" assert await store.search_pages(base_ids=[base["id"]], query="没有关键词重合的查询", embedding_client=SearchEmbedding()) == [] archive.close() restored_archive.close() await engine.dispose() @pytest.mark.asyncio async def test_local_encoder_is_bounded_cpu_and_does_not_need_http(monkeypatch, tmp_path): observed = {} class LocalModel: @staticmethod def list_supported_models(): return [{"model": "BAAI/bge-m3"}] def __init__(self, **kwargs): observed.update(kwargs) def embed(self, texts, *, batch_size): observed["batch_size"] = batch_size return [[1.0, *([0.0] * 1023)] for _ in texts] monkeypatch.setitem(sys.modules, "fastembed", SimpleNamespace(TextEmbedding=LocalModel)) (tmp_path / "onnx").mkdir() (tmp_path / "onnx" / "model.onnx").touch() config = LocalWikiEmbeddingConfig( provider="local", model="bge-embedding-m3", dimensions=1024, batch_size=2, local_threads=2, local_model_path=str(tmp_path), local_process_isolation=False, ) client = StrictWikiEmbeddingClient(config) assert len(await client.embed_query("本地查询")) == 1024 assert observed["model_name"] == "BAAI/bge-m3" assert observed["providers"] == ["CPUExecutionProvider"] assert observed["local_files_only"] is True assert observed["specific_model_path"] == str(tmp_path) assert observed["batch_size"] == observed["threads"] == 2 assert client._client is None assert client.local_encoder_status()["state"] == "ready" assert client.local_encoder_status()["model_files_ready"] is True cold_client = StrictWikiEmbeddingClient(config) status = cold_client.start_local_encoder() assert status["state"] == "starting" assert status["running"] is True await cold_client._local_warmup_task assert cold_client.local_encoder_status()["state"] == "ready" monkeypatch.delenv("DEERFLOW_BGE_M3_MODEL_PATH", raising=False) missing_path = StrictWikiEmbeddingClient( LocalWikiEmbeddingConfig(provider="local", model="bge-embedding-m3", dimensions=1024), ) missing_path.start_local_encoder() await missing_path._local_warmup_task assert missing_path.local_encoder_status()["state"] == "failed" assert "目录未配置" in str(missing_path.local_encoder_status()["last_error"]) @pytest.mark.asyncio async def test_local_encoder_uses_offline_isolated_worker(monkeypatch, tmp_path): (tmp_path / "onnx").mkdir() (tmp_path / "onnx" / "model.onnx").touch() observed: dict = {} class FakeStdout: def __init__(self): self.lines = [json.dumps({"status": "ready"}) + "\n"] def readline(self): return self.lines.pop(0) if self.lines else "" class FakeProcess: def __init__(self): self.stdout = FakeStdout() self.returncode = None self.stdin = SimpleNamespace(write=self.write, flush=lambda: None) def write(self, line): payload = json.loads(line) if payload.get("command") == "shutdown": self.returncode = 0 return self.stdout.lines.append( json.dumps( { "id": payload["id"], "status": "ok", "vectors": [[1.0, *([0.0] * 1023)] for _ in payload["texts"]], } ) + "\n" ) def poll(self): return self.returncode def wait(self, timeout=None): return self.returncode def terminate(self): self.returncode = -15 def kill(self): self.returncode = -9 process = FakeProcess() def fake_popen(command, **kwargs): observed.update({"command": command, **kwargs}) return process monkeypatch.setattr(embedding_module.subprocess, "Popen", fake_popen) client = StrictWikiEmbeddingClient( LocalWikiEmbeddingConfig( provider="local", model="bge-embedding-m3", dimensions=1024, local_model_path=str(tmp_path), ) ) assert len(await client.embed_query("隔离编码")) == 1024 assert "deerflow.integrations.weknora.local_index.embedding_worker" in observed["command"] assert observed["env"]["HF_HUB_OFFLINE"] == "1" assert observed["env"]["PYTHONIOENCODING"] == "utf-8" assert client.local_encoder_status()["process_isolation"] is True await client.close() assert process.poll() == 0 @pytest.mark.asyncio async def test_builtin_skill_without_ownership_record_can_be_distilled(monkeypatch, tmp_path): from starlette.requests import Request from app.gateway.routers import skill_knowledge async def admin(_): return "admin", True class Mappings: async def get_authorized(self, *args, **kwargs): return {"name": "测试库", "weknora_id": "remote"} class Jobs: async def create_job(self, **kwargs): return kwargs app = FastAPI() app.state.config = None app.state.llmwiki_store = Mappings() app.state.skill_knowledge_store = Jobs() # Deliberately no skill_store ownership row: built-ins live on disk. monkeypatch.setattr(skill_knowledge, "_actor", admin) monkeypatch.setattr(skill_knowledge, "_skill_directory", lambda *args: tmp_path) body = skill_knowledge.CreateSyncJobRequest(skill_names=["builtin"], targets=[{"target_type": "weknora", "target_mode": "wiki", "target_id": "mapping"}]) job = await skill_knowledge.create_sync_job(Request({"type": "http", "app": app}), body) assert job["skill_names"] == ["builtin"] @pytest.mark.asyncio async def test_regular_user_can_distill_visible_skill_only_to_writable_wiki(monkeypatch, tmp_path): from starlette.requests import Request from app.gateway.routers import skill_knowledge async def regular_user(_): return "user-1", False class Skills: async def get_visible(self, name, user_id): assert user_id == "user-1" return {"name": name} if name == "visible-skill" else None async def list_favorite_skill_names(self, user_id): assert user_id == "user-1" return [] class Mappings: async def get_authorized(self, mapping_id, user_id, *, write, is_admin): assert (mapping_id, user_id, write, is_admin) == ("mapping", "user-1", True, False) return {"name": "我的 Wiki", "weknora_id": "remote"} class Jobs: async def create_job(self, **kwargs): return kwargs app = FastAPI() app.state.config = None app.state.skill_store = Skills() app.state.llmwiki_store = Mappings() app.state.skill_knowledge_store = Jobs() monkeypatch.setattr(skill_knowledge, "_actor", regular_user) monkeypatch.setattr(skill_knowledge, "_skill_directory", lambda *args: tmp_path) monkeypatch.setattr( skill_knowledge, "get_or_new_skill_storage", lambda **_: SimpleNamespace(get_skills_root_path=lambda: tmp_path), ) body = skill_knowledge.CreateSyncJobRequest( skill_names=["visible-skill"], targets=[{"target_type": "weknora", "target_mode": "wiki", "target_id": "mapping"}], review_mode="required", ) job = await skill_knowledge.create_sync_job(Request({"type": "http", "app": app}), body) assert job["created_by"] == "user-1" assert job["review_mode"] == "off" def test_short_wiki_summary_and_body_still_have_a_vector_section(): sections = chunk_wiki_page({"title": "风扇", "summary": "通过气流帮助散热。", "content_md": "# 风扇\n用于降低设备温度。"}, max_chars=1000, overlap=100, min_chars=80, title_prefix_enabled=True) assert len(sections) == 1 assert "降低设备温度" in sections[0].embedding_text @pytest.mark.asyncio async def test_admin_can_upload_and_activate_offline_encoder(monkeypatch, tmp_path): from starlette.requests import Request from app.gateway.routers import llmwiki_index from deerflow.config.system_settings import SystemSettings async def admin(_): return "admin", True class Embedding: def __init__(self): self.model_path = "" async def reset_local_encoder(self, *, model_path=None): if model_path is not None: self.model_path = model_path def local_encoder_status(self): return { "provider": "local", "model": "bge-embedding-m3", "state": "idle", "model_path": self.model_path, "model_files_ready": bool(self.model_path), } saved = [] monkeypatch.setattr(llmwiki_index, "_actor", admin) monkeypatch.setattr(llmwiki_index, "runtime_home", lambda: tmp_path) monkeypatch.setattr(llmwiki_index, "load_system_settings", SystemSettings) monkeypatch.setattr(llmwiki_index, "save_system_settings", lambda value: saved.append(value) or value) embedding = Embedding() app = FastAPI() app.state.config = SimpleNamespace( llmwiki=SimpleNamespace( local_wiki_index=SimpleNamespace(embedding=SimpleNamespace(provider="local")) ) ) app.state.llmwiki_embedding = embedding request = Request({"type": "http", "headers": [], "app": app}) files = [ UploadFile(filename="model.onnx", file=BytesIO(b"onnx")), UploadFile(filename="config.json", file=BytesIO(b"{}")), ] result = await llmwiki_index.upload_local_encoder( request, files=files, paths=["onnx/model.onnx", "config.json"], ) active = tmp_path / "models" / "bge-m3" assert (active / "onnx" / "model.onnx").read_bytes() == b"onnx" assert result["uploaded_files"] == 2 assert saved[-1].local_wiki_encoder.model_path == str(active.resolve()) assert embedding.model_path == str(active.resolve()) def test_encoder_upload_rejects_parent_traversal(): from app.gateway.routers import llmwiki_index with pytest.raises(Exception) as exc_info: llmwiki_index._safe_relative_upload_path("../model.onnx") assert getattr(exc_info.value, "status_code", None) == 422