514 lines
18 KiB
Python
514 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from types import SimpleNamespace
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from app.gateway import llmwiki_index_scheduler
|
|
from deerflow.config.llmwiki_config import LocalWikiEmbeddingConfig, LocalWikiIndexConfig
|
|
from deerflow.integrations.weknora.local_index.embedding import StrictWikiEmbeddingClient
|
|
from deerflow.integrations.weknora.local_index.normalize import normalize_wiki_page_links
|
|
from deerflow.integrations.weknora.local_index.sync import WikiSyncLeaseLost, WikiSyncService
|
|
from deerflow.integrations.weknora.local_index.vector_cache import WikiVectorCache
|
|
from deerflow.persistence.llmwiki_index.memory import MemoryLlmWikiIndexStore
|
|
|
|
|
|
class FakeEmbedding:
|
|
dimensions = 2
|
|
fingerprint = "fp"
|
|
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
async def embed_texts(self, texts: list[str]) -> list[list[float]]:
|
|
self.calls += 1
|
|
return [[1.0, 0.0] for _ in texts]
|
|
|
|
|
|
class FailingEmbedding(FakeEmbedding):
|
|
async def embed_texts(self, texts: list[str]) -> list[list[float]]:
|
|
raise RuntimeError("embedding unavailable")
|
|
|
|
|
|
class LeaseStealingEmbedding(FakeEmbedding):
|
|
def __init__(self, store: MemoryLlmWikiIndexStore) -> None:
|
|
super().__init__()
|
|
self.store = store
|
|
|
|
async def embed_texts(self, texts: list[str]) -> list[list[float]]:
|
|
state = self.store.states["kb"]
|
|
state["lease_owner"] = "replacement-worker"
|
|
state["active_sync_id"] = "replacement-sync"
|
|
return await super().embed_texts(texts)
|
|
|
|
|
|
class LeaseStealingFailingEmbedding(LeaseStealingEmbedding):
|
|
async def embed_texts(self, texts: list[str]) -> list[list[float]]:
|
|
await super().embed_texts(texts)
|
|
raise RuntimeError("embedding failed after lease takeover")
|
|
|
|
|
|
class NewFingerprintEmbedding(FakeEmbedding):
|
|
fingerprint = "fp-new"
|
|
|
|
|
|
class FakeWeKnora:
|
|
def __init__(self) -> None:
|
|
self.pages = {
|
|
"one": {
|
|
"id": "p1",
|
|
"slug": "one",
|
|
"title": "One",
|
|
"summary": "A long enough summary for indexing and retrieval.",
|
|
"content": "# Intro\nA long enough Wiki paragraph for indexing and semantic retrieval.",
|
|
"status": "published",
|
|
"version": 1,
|
|
},
|
|
}
|
|
|
|
async def get_knowledge_base(self, knowledge_base_id: str):
|
|
return {"id": knowledge_base_id, "indexing_strategy": {"wiki_enabled": True}}
|
|
|
|
async def list_wiki_pages(self, knowledge_base_id: str, *, page: int, page_size: int):
|
|
return {
|
|
"pages": [{"slug": item["slug"], "title": item["title"]} for item in self.pages.values()],
|
|
"total": len(self.pages),
|
|
"total_pages": 1,
|
|
}
|
|
|
|
async def get_wiki_page(self, knowledge_base_id: str, slug: str):
|
|
return dict(self.pages[slug])
|
|
|
|
|
|
class LinkedFakeWeKnora(FakeWeKnora):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.pages["one"].update(
|
|
{
|
|
"out_links": [{"slug": "two"}],
|
|
"content": "# Intro\nOne links to [[Two Page|the second page]].",
|
|
}
|
|
)
|
|
self.pages["two"] = {
|
|
"id": "p2",
|
|
"slug": "two",
|
|
"title": "Two Page",
|
|
"summary": "A second page with enough text for the local index.",
|
|
"content": "# Details\nThe second page contains connected evidence.",
|
|
"inLinks": ["one"],
|
|
"status": "published",
|
|
"version": 1,
|
|
}
|
|
|
|
async def get_wiki_graph(self, knowledge_base_id: str, *, limit: int):
|
|
return {
|
|
"nodes": [
|
|
{"id": "node-one", "slug": "one", "title": "One"},
|
|
{"id": "node-two", "slug": "two", "title": "Two Page"},
|
|
],
|
|
"edges": [{"source": "node-one", "target": "node-two"}],
|
|
}
|
|
|
|
|
|
class LeaseStealingGraphWeKnora(FakeWeKnora):
|
|
def __init__(self, store: MemoryLlmWikiIndexStore) -> None:
|
|
super().__init__()
|
|
self.store = store
|
|
|
|
async def get_wiki_graph(self, knowledge_base_id: str, *, limit: int):
|
|
del knowledge_base_id, limit
|
|
state = self.store.states["kb"]
|
|
state["lease_owner"] = "replacement-worker"
|
|
state["active_sync_id"] = "replacement-sync"
|
|
return {"nodes": [], "edges": []}
|
|
|
|
|
|
def test_normalize_wiki_page_links_handles_provider_shapes_and_markdown() -> None:
|
|
links = normalize_wiki_page_links(
|
|
{
|
|
"inLinks": [{"page_slug": "source/a#overview"}],
|
|
"out_links": ["target/b", {"path": "target/c"}],
|
|
"content": "See [[Target D#details|details]] and [[target/b]].",
|
|
}
|
|
)
|
|
|
|
assert links == {
|
|
"in_links": ["source/a"],
|
|
"out_links": ["Target D", "target/b", "target/c"],
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_configured_dimensions_are_validated_but_not_sent_to_bge() -> None:
|
|
requests: list[dict] = []
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
payload = json.loads(request.content)
|
|
requests.append(payload)
|
|
if "dimensions" in payload:
|
|
return httpx.Response(400, json={"detail": "dimensions is not supported"})
|
|
return httpx.Response(
|
|
200,
|
|
json={"data": [{"index": 0, "embedding": [1.0, 0.0, 0.0]}]},
|
|
)
|
|
|
|
config = LocalWikiEmbeddingConfig(
|
|
model="bge-embedding-m3",
|
|
base_url="http://embedding.test/v1",
|
|
dimensions=3,
|
|
max_retries=0,
|
|
)
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client:
|
|
client = StrictWikiEmbeddingClient(config, http_client=http_client)
|
|
vector = await client.embed_query("Wiki page")
|
|
|
|
assert vector == [1.0, 0.0, 0.0]
|
|
assert requests == [{"model": "bge-embedding-m3", "input": ["Wiki page"]}]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_embedding_error_includes_upstream_reason(caplog) -> None:
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(400, json={"detail": "unsupported request parameter: dimensions"})
|
|
|
|
config = LocalWikiEmbeddingConfig(
|
|
model="bge-embedding-m3",
|
|
base_url="http://embedding.test/v1",
|
|
max_retries=0,
|
|
)
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client:
|
|
client = StrictWikiEmbeddingClient(config, http_client=http_client)
|
|
with pytest.raises(RuntimeError, match="unsupported request parameter: dimensions"):
|
|
await client.embed_query("Wiki page")
|
|
|
|
assert "unsupported request parameter: dimensions" in caplog.text
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_omitted_dimensions_are_discovered_from_first_response() -> None:
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(
|
|
200,
|
|
json={"data": [{"index": 0, "embedding": [0.0, 1.0]}]},
|
|
)
|
|
|
|
config = LocalWikiEmbeddingConfig(
|
|
model="bge-embedding-m3",
|
|
base_url="http://embedding.test/v1",
|
|
dimensions=None,
|
|
max_retries=0,
|
|
)
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client:
|
|
client = StrictWikiEmbeddingClient(config, http_client=http_client)
|
|
await client.embed_query("dimension probe")
|
|
|
|
assert client.dimensions == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sync_is_incremental_and_full_scan_deletes_missing_pages() -> None:
|
|
config = LocalWikiIndexConfig(chunk_min_chars=10, min_similarity=0)
|
|
store = MemoryLlmWikiIndexStore()
|
|
embedding = FakeEmbedding()
|
|
client = FakeWeKnora()
|
|
service = WikiSyncService(config, store, embedding, WikiVectorCache(store), client)
|
|
mapping = {"id": "kb", "weknora_id": "remote", "name": "Wiki"}
|
|
|
|
first = await service.sync_mapping(mapping)
|
|
assert first["changed_pages"] == 1
|
|
assert embedding.calls == 1
|
|
statuses = await store.list_page_index_status("kb")
|
|
assert statuses[0]["title"] == "One"
|
|
assert statuses[0]["vector_count"] > 0
|
|
assert store.states["kb"]["remote_page_count"] == 1
|
|
second = await service.sync_mapping(mapping)
|
|
assert second["changed_pages"] == 0
|
|
assert embedding.calls == 1
|
|
|
|
client.pages.clear()
|
|
deleted = await service.sync_mapping(mapping)
|
|
assert deleted["deleted_pages"] == 1
|
|
assert next(iter(store.pages.values()))["is_remote_deleted"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sync_persists_mapping_scoped_wiki_edges_for_adjacency() -> None:
|
|
config = LocalWikiIndexConfig(chunk_min_chars=10, min_similarity=0)
|
|
store = MemoryLlmWikiIndexStore()
|
|
client = LinkedFakeWeKnora()
|
|
mapping = {"id": "kb", "weknora_id": "remote", "name": "Wiki"}
|
|
|
|
result = await WikiSyncService(config, store, FakeEmbedding(), WikiVectorCache(store), client).sync_mapping(mapping)
|
|
|
|
assert result["graph_edges"] == 1
|
|
assert result["graph_changed"] is True
|
|
neighbors = await store.get_adjacent_pages(["kb"], {"kb": ["one"]}, direction="outgoing")
|
|
assert len(neighbors) == 1
|
|
assert neighbors[0] == {
|
|
**neighbors[0],
|
|
"knowledge_base_mapping_id": "kb",
|
|
"seed_slug": "one",
|
|
"neighbor_slug": "two",
|
|
"direction": "outgoing",
|
|
"source_slug": "one",
|
|
"target_slug": "two",
|
|
}
|
|
assert neighbors[0]["page"]["title"] == "Two Page"
|
|
|
|
# A seed belonging to another mapping cannot traverse this mapping's edge.
|
|
assert await store.get_adjacent_pages(["other"], {"other": ["one"]}, direction="both") == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sync_losing_lease_during_graph_fetch_cannot_delete_or_replace_snapshot() -> None:
|
|
config = LocalWikiIndexConfig(chunk_min_chars=10, min_similarity=0)
|
|
store = MemoryLlmWikiIndexStore()
|
|
mapping = {"id": "kb", "weknora_id": "remote", "name": "Wiki"}
|
|
initial_client = LinkedFakeWeKnora()
|
|
await WikiSyncService(
|
|
config,
|
|
store,
|
|
FakeEmbedding(),
|
|
WikiVectorCache(store),
|
|
initial_client,
|
|
).sync_mapping(mapping)
|
|
page_snapshot = {
|
|
page_id: dict(page)
|
|
for page_id, page in store.pages.items()
|
|
}
|
|
edge_snapshot = {
|
|
mapping_id: [dict(edge) for edge in edges]
|
|
for mapping_id, edges in store.edges.items()
|
|
}
|
|
|
|
stealing_client = LeaseStealingGraphWeKnora(store)
|
|
stealing_client.pages.clear()
|
|
with pytest.raises(WikiSyncLeaseLost, match="lease was lost"):
|
|
await WikiSyncService(
|
|
config,
|
|
store,
|
|
FakeEmbedding(),
|
|
WikiVectorCache(store),
|
|
stealing_client,
|
|
).sync_mapping(mapping)
|
|
|
|
assert store.pages == page_snapshot
|
|
assert store.edges == edge_snapshot
|
|
assert store.states["kb"]["lease_owner"] == "replacement-worker"
|
|
assert store.states["kb"]["active_sync_id"] == "replacement-sync"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sync_losing_lease_during_embedding_cannot_replace_page_snapshot() -> None:
|
|
config = LocalWikiIndexConfig(chunk_min_chars=10, min_similarity=0)
|
|
store = MemoryLlmWikiIndexStore()
|
|
client = FakeWeKnora()
|
|
mapping = {"id": "kb", "weknora_id": "remote", "name": "Wiki"}
|
|
await WikiSyncService(
|
|
config,
|
|
store,
|
|
FakeEmbedding(),
|
|
WikiVectorCache(store),
|
|
client,
|
|
).sync_mapping(mapping)
|
|
page_snapshot = {
|
|
page_id: dict(page)
|
|
for page_id, page in store.pages.items()
|
|
}
|
|
vector_snapshot = {
|
|
page_id: [dict(vector) for vector in vectors]
|
|
for page_id, vectors in store.vectors.items()
|
|
}
|
|
client.pages["one"]["content"] += " Changed by a stale scan."
|
|
|
|
with pytest.raises(WikiSyncLeaseLost, match="lease was lost"):
|
|
await WikiSyncService(
|
|
config,
|
|
store,
|
|
LeaseStealingEmbedding(store),
|
|
WikiVectorCache(store),
|
|
client,
|
|
).sync_mapping(mapping)
|
|
|
|
assert store.pages == page_snapshot
|
|
assert store.vectors == vector_snapshot
|
|
assert store.states["kb"]["lease_owner"] == "replacement-worker"
|
|
assert store.states["kb"]["active_sync_id"] == "replacement-sync"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_page_failure_fencing_propagates_lease_loss_without_wrapping() -> None:
|
|
config = LocalWikiIndexConfig(chunk_min_chars=10, min_similarity=0)
|
|
store = MemoryLlmWikiIndexStore()
|
|
client = FakeWeKnora()
|
|
mapping = {"id": "kb", "weknora_id": "remote", "name": "Wiki"}
|
|
|
|
with pytest.raises(WikiSyncLeaseLost, match="recording a page failure"):
|
|
await WikiSyncService(
|
|
config,
|
|
store,
|
|
LeaseStealingFailingEmbedding(store),
|
|
WikiVectorCache(store),
|
|
client,
|
|
).sync_mapping(mapping)
|
|
|
|
assert store.pages == {}
|
|
assert store.states["kb"]["lease_owner"] == "replacement-worker"
|
|
assert store.states["kb"]["active_sync_id"] == "replacement-sync"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_page_failure_keeps_previous_generation_and_marks_partial() -> None:
|
|
config = LocalWikiIndexConfig(chunk_min_chars=10, min_similarity=0)
|
|
store = MemoryLlmWikiIndexStore()
|
|
client = FakeWeKnora()
|
|
mapping = {"id": "kb", "weknora_id": "remote", "name": "Wiki"}
|
|
await WikiSyncService(config, store, FakeEmbedding(), WikiVectorCache(store), client).sync_mapping(mapping)
|
|
page = next(iter(store.pages.values()))
|
|
old_hash = page["content_hash"]
|
|
old_vectors = list(store.vectors[page["id"]])
|
|
client.pages["one"]["content"] += " Changed."
|
|
|
|
result = await WikiSyncService(config, store, FailingEmbedding(), WikiVectorCache(store), client).sync_mapping(mapping)
|
|
|
|
assert result["status"] == "partial"
|
|
assert store.states["kb"]["state"] == "partial"
|
|
assert store.states["kb"]["failed_page_count"] == 1
|
|
assert store.pages[page["id"]]["content_hash"] == old_hash
|
|
assert store.pages[page["id"]]["index_status"] == "ready"
|
|
assert store.vectors[page["id"]] == old_vectors
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fingerprint_change_automatically_rebuilds_entire_mapping() -> None:
|
|
config = LocalWikiIndexConfig(chunk_min_chars=10, min_similarity=0)
|
|
store = MemoryLlmWikiIndexStore()
|
|
client = FakeWeKnora()
|
|
mapping = {"id": "kb", "weknora_id": "remote", "name": "Wiki"}
|
|
await WikiSyncService(config, store, FakeEmbedding(), WikiVectorCache(store), client).sync_mapping(mapping)
|
|
|
|
embedding = NewFingerprintEmbedding()
|
|
result = await WikiSyncService(config, store, embedding, WikiVectorCache(store), client).sync_mapping(mapping)
|
|
|
|
assert result["status"] == "completed"
|
|
assert result["changed_pages"] == 1
|
|
assert embedding.calls == 1
|
|
assert store.states["kb"]["state"] == "idle"
|
|
assert store.states["kb"]["embedding_fingerprint"] == "fp-new"
|
|
assert {vector["embedding_fingerprint"] for vectors in store.vectors.values() for vector in vectors} == {"fp-new"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_auto_sync_scheduler_scans_immediately_then_only_when_due(monkeypatch) -> None:
|
|
now = 0.0
|
|
refreshes: list[float] = []
|
|
syncs: list[float] = []
|
|
sleeps: list[float] = []
|
|
config = SimpleNamespace(
|
|
enabled=True,
|
|
auto_sync=True,
|
|
sync_interval_seconds=25,
|
|
)
|
|
|
|
async def fake_refresh(_app):
|
|
refreshes.append(now)
|
|
return config
|
|
|
|
async def fake_sleep(seconds: float) -> None:
|
|
nonlocal now
|
|
sleeps.append(seconds)
|
|
if len(syncs) >= 2:
|
|
raise asyncio.CancelledError
|
|
now += seconds
|
|
|
|
async def fake_sync_once(_app, *, config=None) -> None:
|
|
syncs.append(now)
|
|
|
|
monkeypatch.setattr(llmwiki_index_scheduler, "monotonic", lambda: now)
|
|
monkeypatch.setattr(llmwiki_index_scheduler.asyncio, "sleep", fake_sleep)
|
|
monkeypatch.setattr(llmwiki_index_scheduler, "_sync_once", fake_sync_once)
|
|
monkeypatch.setattr(
|
|
llmwiki_index_scheduler,
|
|
"_refresh_local_wiki_config",
|
|
fake_refresh,
|
|
)
|
|
app = SimpleNamespace(
|
|
state=SimpleNamespace(
|
|
llmwiki_sync_service=object(),
|
|
llmwiki_store=object(),
|
|
)
|
|
)
|
|
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await llmwiki_index_scheduler.local_wiki_index_scheduler_loop(app)
|
|
|
|
assert refreshes == [0.0, 10.0, 20.0, 30.0]
|
|
assert syncs == [0.0, 30.0]
|
|
assert sleeps == [10.0, 10.0, 10.0, 10.0]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_auto_sync_scheduler_refreshes_provider_between_syncs(monkeypatch) -> None:
|
|
now = 0.0
|
|
seen_providers: list[str] = []
|
|
syncs: list[str] = []
|
|
providers = iter(("openai", "local", "local"))
|
|
|
|
async def fake_refresh(_app):
|
|
provider = next(providers)
|
|
seen_providers.append(provider)
|
|
_app.state.llmwiki_runtime_signature = provider
|
|
return SimpleNamespace(
|
|
enabled=True,
|
|
auto_sync=True,
|
|
sync_interval_seconds=120,
|
|
embedding=SimpleNamespace(provider=provider),
|
|
)
|
|
|
|
async def fake_sync_once(_app, *, config=None) -> None:
|
|
syncs.append(config.embedding.provider)
|
|
|
|
async def fake_sleep(seconds: float) -> None:
|
|
nonlocal now
|
|
now += seconds
|
|
if len(seen_providers) >= 3:
|
|
raise asyncio.CancelledError
|
|
|
|
monkeypatch.setattr(llmwiki_index_scheduler, "monotonic", lambda: now)
|
|
monkeypatch.setattr(llmwiki_index_scheduler.asyncio, "sleep", fake_sleep)
|
|
monkeypatch.setattr(llmwiki_index_scheduler, "_sync_once", fake_sync_once)
|
|
monkeypatch.setattr(
|
|
llmwiki_index_scheduler,
|
|
"_refresh_local_wiki_config",
|
|
fake_refresh,
|
|
)
|
|
app = SimpleNamespace(
|
|
state=SimpleNamespace(
|
|
llmwiki_sync_service=object(),
|
|
llmwiki_store=object(),
|
|
)
|
|
)
|
|
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await llmwiki_index_scheduler.local_wiki_index_scheduler_loop(app)
|
|
|
|
assert seen_providers == ["openai", "local", "local"]
|
|
assert syncs == ["openai", "local"]
|
|
|
|
|
|
def test_local_wiki_index_defaults_to_automatic_sync() -> None:
|
|
config = LocalWikiIndexConfig()
|
|
|
|
assert config.auto_sync is True
|
|
assert config.sync_interval_seconds == 30
|
|
|
|
|
|
def test_local_wiki_index_accepts_legacy_sync_flag() -> None:
|
|
config = LocalWikiIndexConfig.model_validate({"sync_enabled": True})
|
|
|
|
assert config.auto_sync is True
|
|
assert config.sync_enabled is True
|