From 62de7f3425f748f5cd43264340041aafef50635e Mon Sep 17 00:00:00 2001 From: Cordis798 <2378285593@qq.com> Date: Fri, 30 Oct 2026 07:08:01 +0800 Subject: [PATCH] perf(knowledge): bound concurrent WeKnora document checks (#7707) Co-authored-by: Cordis798 <284918518+Cordis798@users.noreply.github.com> --- .../test_weknora_document_concurrency.py | 170 ++++++++++++++++++ examples/deerflow-extension-weknora/AGENTS.md | 4 + examples/deerflow-extension-weknora/README.md | 6 + .../deerflow_extension_weknora/provider.py | 36 +++- 4 files changed, 209 insertions(+), 7 deletions(-) create mode 100644 backend/tests/test_weknora_document_concurrency.py diff --git a/backend/tests/test_weknora_document_concurrency.py b/backend/tests/test_weknora_document_concurrency.py new file mode 100644 index 0000000000..58e0368849 --- /dev/null +++ b/backend/tests/test_weknora_document_concurrency.py @@ -0,0 +1,170 @@ +"""Bounded, request-local detail checks retain WeKnora's authorization gate.""" + +import asyncio +from collections import Counter +from pathlib import Path + +import httpx +import pytest + +from deerflow.knowledge import KnowledgeRequestContext + + +@pytest.fixture +def setup_provider(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "examples/deerflow-extension-weknora")) + from deerflow_extension_weknora import WeKnoraProvider + + monkeypatch.setenv("WK_CONCURRENCY_KEY", "synthetic") + + def create(*, gate=False, failure=None, hold=False): + state = {"active": 0, "peak": 0, "calls": Counter(), "searches": 0} + started, release = asyncio.Event(), asyncio.Event() + + async def handle(request): + assert request.headers["X-API-Key"] == "synthetic" + assert request.headers["X-Tenant-ID"] == "1" + path = request.url.path + if path == "/api/v1/knowledge-bases": + data = [{"id": "kb", "name": "Synthetic"}] + elif path.endswith("hybrid-search"): + state["searches"] += 1 + data = [{"knowledge_id": f"doc-{i}", "content": f"chunk {i}"} for i in range(8)] + data.insert(1, {"knowledge_id": "doc-0", "content": "second chunk"}) + else: + doc = path.rsplit("/", 1)[1] + state["calls"][doc] += 1 + state["active"] += 1 + state["peak"] = max(state["peak"], state["active"]) + started.set() + try: + if state["active"] == 4: + release.set() + if gate: + await asyncio.wait_for(release.wait(), 2) + if failure and doc == "doc-0": + return httpx.Response(failure, json={"success": False}) + if hold: + await asyncio.Event().wait() + await asyncio.sleep(0) + data = {"id": doc, "knowledge_base_id": "kb", "parse_status": "completed", "enable_status": "enabled", "title": doc} + finally: + state["active"] -= 1 + return httpx.Response(200, json={"success": True, "data": data}) + + provider = WeKnoraProvider({"spaces": [{"base_url": "https://weknora.invalid", "tenant_id": "1", "api_key_env": "WK_CONCURRENCY_KEY", "user_ids": ["alice"]}]}, transport=httpx.MockTransport(handle)) + return provider, state, started + + return create + + +@pytest.mark.asyncio +@pytest.mark.parametrize("selected", [False, True]) +async def test_details_overlap_with_a_bound_and_preserve_hit_order(setup_provider, selected): + provider, state, _ = setup_provider(gate=True) + ctx = KnowledgeRequestContext("alice", "lead_agent") + if selected: + catalog = await provider.catalog(context=ctx, page=1, page_size=20, search="") + dataset = catalog.items[0].id + ctx = KnowledgeRequestContext("alice", "lead_agent", {"version": 1, "mode": "selected", "dataset_ids": [dataset], "document_filters": [{"dataset_id": dataset, "document_ids": [f"doc-{i}" for i in range(8)]}]}) + result = await provider.search_scoped("q", context=ctx) + assert state["peak"] == 4 + assert state["active"] == 0 + assert state["calls"] == Counter({f"doc-{i}": 1 for i in range(8)}) + assert [h.content for h in result.hits] == ["chunk 0", "second chunk", *[f"chunk {i}" for i in range(1, 8)]] + assert [h.source.title for h in result.hits] == ["doc-0", "doc-0", *[f"doc-{i}" for i in range(1, 8)]] + # A second call must revalidate documents rather than reuse authorization. + await provider.search_scoped("q", context=ctx) + assert all(n == 2 for n in state["calls"].values()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure,error", [(403, PermissionError), (404, FileNotFoundError), (500, RuntimeError)]) +async def test_failure_drains_other_detail_requests_before_returning(setup_provider, failure, error): + provider, state, _ = setup_provider(gate=True, failure=failure, hold=True) + with pytest.raises(error): + await asyncio.wait_for(provider.search_scoped("q", context=KnowledgeRequestContext("alice")), 5) + assert state["peak"] == 4 + assert state["active"] == 0 + assert len(state["calls"]) == 4 # Later batches never start after failure. + + +@pytest.mark.asyncio +async def test_cancellation_drains_detail_requests(setup_provider): + provider, state, started = setup_provider(hold=True) + task = asyncio.create_task(provider.search_scoped("q", context=KnowledgeRequestContext("alice"))) + await asyncio.wait_for(started.wait(), 5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert state["active"] == 0 + + +@pytest.mark.asyncio +async def test_selected_document_denial_prevents_search(setup_provider): + provider, state, _ = setup_provider(gate=True, failure=403, hold=True) + catalog = await provider.catalog(context=KnowledgeRequestContext("alice"), page=1, page_size=20, search="") + dataset = catalog.items[0].id + ctx = KnowledgeRequestContext("alice", "lead_agent", {"version": 1, "mode": "selected", "dataset_ids": [dataset], "document_filters": [{"dataset_id": dataset, "document_ids": [f"doc-{i}" for i in range(8)]}]}) + with pytest.raises(PermissionError): + await asyncio.wait_for(provider.search_scoped("q", context=ctx), 5) + assert state["searches"] == 0 + assert state["active"] == 0 + + +@pytest.mark.asyncio +async def test_equal_native_ids_in_different_tenants_are_revalidated(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "examples/deerflow-extension-weknora")) + from deerflow_extension_weknora import WeKnoraProvider + + monkeypatch.setenv("WK_ONE", "key-one") + monkeypatch.setenv("WK_TWO", "key-two") + details = [] + + async def handle(request): + tenant = request.headers["X-Tenant-ID"] + assert request.headers["X-API-Key"] == {"1": "key-one", "2": "key-two"}[tenant] + if request.url.path == "/api/v1/knowledge-bases": + data = [{"id": "kb", "name": tenant}] + elif request.url.path.endswith("hybrid-search"): + data = [{"knowledge_id": "same-id", "content": f"evidence {tenant}"}] + else: + details.append(tenant) + data = {"id": "same-id", "knowledge_base_id": "kb", "parse_status": "completed", "enable_status": "enabled", "title": tenant} + return httpx.Response(200, json={"success": True, "data": data}) + + provider = WeKnoraProvider( + {"spaces": [{"base_url": "https://weknora.invalid", "tenant_id": tenant, "api_key_env": key, "user_ids": ["alice"]} for tenant, key in [("1", "WK_ONE"), ("2", "WK_TWO")]]}, transport=httpx.MockTransport(handle) + ) + result = await provider.search_scoped("q", context=KnowledgeRequestContext("alice")) + assert details == ["1", "2"] + assert [hit.source.title for hit in result.hits] == ["1", "2"] + + +@pytest.mark.asyncio +async def test_result_cap_checks_empty_hits_but_not_the_unconsumed_suffix(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2] / "examples/deerflow-extension-weknora")) + from deerflow_extension_weknora import WeKnoraProvider + + monkeypatch.setenv("WK_CAP_KEY", "synthetic") + details = [] + + async def handle(request): + path = request.url.path + if path == "/api/v1/knowledge-bases": + data = [{"id": "first", "name": "First"}, {"id": "last", "name": "Last"}] + elif path.endswith("first/hybrid-search"): + data = [{"knowledge_id": "first-doc", "content": f"chunk {i}"} for i in range(99)] + elif path.endswith("last/hybrid-search"): + data = [{"knowledge_id": "empty-doc", "content": ""}, {"knowledge_id": "last-doc", "content": "last chunk"}, {"knowledge_id": "unsafe/path", "content": "unconsumed"}] + else: + doc = path.rsplit("/", 1)[1] + details.append(doc) + data = {"id": doc, "knowledge_base_id": "first" if doc == "first-doc" else "last", "parse_status": "completed", "enable_status": "enabled", "title": doc} + return httpx.Response(200, json={"success": True, "data": data}) + + provider = WeKnoraProvider({"spaces": [{"base_url": "https://weknora.invalid", "tenant_id": "1", "api_key_env": "WK_CAP_KEY", "user_ids": ["alice"]}]}, transport=httpx.MockTransport(handle)) + result = await provider.search_scoped("q", context=KnowledgeRequestContext("alice")) + assert len(result.hits) == 100 + assert result.hits[-1].content == "last chunk" + assert details == ["first-doc", "empty-doc", "last-doc"] diff --git a/examples/deerflow-extension-weknora/AGENTS.md b/examples/deerflow-extension-weknora/AGENTS.md index 03e36d1588..821a63bbd6 100644 --- a/examples/deerflow-extension-weknora/AGENTS.md +++ b/examples/deerflow-extension-weknora/AGENTS.md @@ -28,3 +28,7 @@ Wiki lexical search retries only successful empty pages with up to four distinct whitespace-separated keywords of at least two characters. Keep retries in the same authorized base, validate every response, deduplicate title/excerpt pairs, and never retry provider/auth/schema errors. Nonempty phrase results stay intact. + +Batch document checks at four per search; retain dataset/ID request-local keys +and hit order. Drain sibling tasks on errors/cancellation. Validate selected +documents before search, and only the consumed hit prefix before returning it. diff --git a/examples/deerflow-extension-weknora/README.md b/examples/deerflow-extension-weknora/README.md index a0926517a8..55c34d216a 100644 --- a/examples/deerflow-extension-weknora/README.md +++ b/examples/deerflow-extension-weknora/README.md @@ -220,3 +220,9 @@ at most four distinct whitespace-separated keywords (at least two characters), within the same authorized knowledge base, and merges duplicate evidence. This helps model-generated keyword combinations match the upstream lexical search; it does not perform unrestricted search or retry authorization/provider errors. + +Document details are checked in batches of at most four concurrent requests per +search. Every selected document is validated before retrieval; retrieved evidence +is validated before returning results. Checks remain request-local and results +retain provider order. A failed or cancelled batch cancels and drains its other +requests; it never falls back to broader scope or returns partial evidence. diff --git a/examples/deerflow-extension-weknora/deerflow_extension_weknora/provider.py b/examples/deerflow-extension-weknora/deerflow_extension_weknora/provider.py index 001f8f8802..9169bf8c59 100644 --- a/examples/deerflow-extension-weknora/deerflow_extension_weknora/provider.py +++ b/examples/deerflow-extension-weknora/deerflow_extension_weknora/provider.py @@ -21,6 +21,7 @@ from deerflow.knowledge_scope import execution_scope from .config import Settings, Space, native_id _MAX_RESPONSE_BYTES = 4 * 1024 * 1024 +_DOCUMENT_CONCURRENCY = 4 @dataclass(frozen=True) @@ -181,6 +182,23 @@ class WeKnoraProvider: async def search(self, query: str) -> RetrievalResult: raise PermissionError("WeKnora retrieval requires a trusted KnowledgeRequestContext") + async def _load_documents(self, ds: _Dataset, ids: list[str], documents: dict[tuple[str, str], dict[str, Any]]) -> None: + # Cache only within this search, keyed by authorized dataset and native ID. + missing = list(dict.fromkeys(document_id for document_id in ids if (ds.id, document_id) not in documents)) + for start in range(0, len(missing), _DOCUMENT_CONCURRENCY): + batch = missing[start : start + _DOCUMENT_CONCURRENCY] + tasks = [asyncio.create_task(self._document(ds, document_id)) for document_id in batch] + try: + values = await asyncio.gather(*tasks) + except BaseException: + # gather alone leaves siblings running after an ordinary error. + # Drain them on failure/cancellation, preserving the original error. + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise + documents.update(((ds.id, document_id), value) for document_id, value in zip(batch, values, strict=True)) + async def search_scoped(self, query: str, *, context: KnowledgeRequestContext) -> RetrievalResult: scope = execution_scope(context.scope) if context.scope is not None else {"mode": "all"} if scope["mode"] == "disabled": @@ -199,8 +217,7 @@ class WeKnoraProvider: # Validate every constrained resource before any search; no partial # success and no unfiltered retry if any selection fails. for dataset_id, ids in filters.items(): - for document_id in ids: - documents[dataset_id, document_id] = await self._document(datasets[dataset_id], document_id) + await self._load_documents(datasets[dataset_id], ids, documents) hits = [] for dataset_id in selected: if len(hits) >= 100: @@ -217,20 +234,25 @@ class WeKnoraProvider: results = _items(payload["data"]) if len(results) > 100: raise RuntimeError("WeKnora response exceeds the supported result limit") + candidates = [] + remaining = 100 - len(hits) for item in results: - if len(hits) >= 100: + if remaining == 0: break document_id = native_id(item.get("knowledge_id")) if dataset_id in filters and document_id not in filters[dataset_id]: raise PermissionError("WeKnora returned evidence outside the selected documents") - cache_key = (dataset_id, document_id) - if cache_key not in documents: - documents[cache_key] = await self._document(ds, document_id) + text = ds.binding.text(item.get("content"), 100_000) + candidates.append((item, document_id, text)) + if text: + remaining -= 1 + await self._load_documents(ds, [document_id for _, document_id, _ in candidates], documents) + for item, document_id, text in candidates: # SearchResult has no knowledge_base_id; the authoritative detail # endpoint above verifies membership even for unfiltered search. - text = ds.binding.text(item.get("content"), 100_000) if not text: continue + cache_key = (dataset_id, document_id) title = ds.binding.text(item.get("knowledge_title") or item.get("knowledge_filename") or documents[cache_key].get("title") or documents[cache_key].get("file_name"), 500) source = KnowledgeSource(provider="weknora", title=title, excerpt=text, dataset_name=ds.name) if title else None # v0.8.3 gives offsets/chunk indexes, not a verified page or public