perf(knowledge): bound concurrent WeKnora document checks (#7707)

Co-authored-by: Cordis798 <284918518+Cordis798@users.noreply.github.com>

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"]

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.

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.

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