284 lines
9.5 KiB
Python
284 lines
9.5 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from pathlib import Path
|
|
from threading import Event
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import yaml
|
|
from fastapi import FastAPI
|
|
from starlette.requests import Request
|
|
|
|
from app.gateway import llmwiki_index_runtime
|
|
from app.gateway.routers import llmwiki_index
|
|
from deerflow.config.llmwiki_config import LocalWikiEmbeddingConfig, LocalWikiIndexConfig
|
|
from deerflow.config.system_settings import LocalWikiEncoderSettings, SystemSettings
|
|
from deerflow.integrations.weknora.local_index.embedding import StrictWikiEmbeddingClient, WikiEmbeddingError
|
|
|
|
|
|
def _source_config(embedding: LocalWikiEmbeddingConfig):
|
|
return SimpleNamespace(
|
|
llmwiki=SimpleNamespace(
|
|
local_wiki_index=LocalWikiIndexConfig(enabled=True, embedding=embedding)
|
|
)
|
|
)
|
|
|
|
|
|
def test_runtime_mode_can_switch_remote_config_to_uploaded_local_model(tmp_path: Path) -> None:
|
|
configured = LocalWikiEmbeddingConfig(
|
|
provider="openai_compatible",
|
|
model="remote-embedding",
|
|
base_url="https://embedding.example.test/v1",
|
|
dimensions=1536,
|
|
)
|
|
source = _source_config(configured)
|
|
|
|
normal = llmwiki_index_runtime.resolve_effective_local_wiki_index_config(
|
|
source,
|
|
LocalWikiEncoderSettings(mode="configured"),
|
|
)
|
|
assert normal.embedding.provider == "openai_compatible"
|
|
assert normal.embedding.model == "remote-embedding"
|
|
assert normal.embedding.dimensions == 1536
|
|
|
|
local = llmwiki_index_runtime.resolve_effective_local_wiki_index_config(
|
|
source,
|
|
LocalWikiEncoderSettings(
|
|
mode="local",
|
|
model_path=str(tmp_path),
|
|
model_revision="revision-2",
|
|
),
|
|
)
|
|
assert local.embedding.provider == "local"
|
|
assert local.embedding.model == "bge-embedding-m3"
|
|
assert local.embedding.dimensions == 1024
|
|
assert local.embedding.local_model_path == str(tmp_path)
|
|
assert local.embedding.local_model_revision == "revision-2"
|
|
assert local.embedding.base_url == ""
|
|
assert local.embedding.api_key == ""
|
|
|
|
|
|
def test_local_model_revision_changes_embedding_fingerprint(tmp_path: Path) -> None:
|
|
common = {
|
|
"provider": "local",
|
|
"model": "bge-embedding-m3",
|
|
"dimensions": 1024,
|
|
"local_model_path": str(tmp_path),
|
|
}
|
|
first = StrictWikiEmbeddingClient(
|
|
LocalWikiEmbeddingConfig(**common, local_model_revision="first")
|
|
)
|
|
second = StrictWikiEmbeddingClient(
|
|
LocalWikiEmbeddingConfig(**common, local_model_revision="second")
|
|
)
|
|
assert first.fingerprint != second.fingerprint
|
|
|
|
|
|
def test_local_model_path_changes_embedding_fingerprint(tmp_path: Path) -> None:
|
|
first_path = tmp_path / "revision-a"
|
|
second_path = tmp_path / "revision-b"
|
|
common = {
|
|
"provider": "local",
|
|
"model": "bge-embedding-m3",
|
|
"dimensions": 1024,
|
|
"local_model_revision": "same-revision",
|
|
}
|
|
first = StrictWikiEmbeddingClient(
|
|
LocalWikiEmbeddingConfig(**common, local_model_path=str(first_path))
|
|
)
|
|
second = StrictWikiEmbeddingClient(
|
|
LocalWikiEmbeddingConfig(**common, local_model_path=str(second_path))
|
|
)
|
|
assert first.fingerprint != second.fingerprint
|
|
|
|
|
|
def test_uploaded_local_mode_always_uses_bge_m3_contract(tmp_path: Path) -> None:
|
|
configured = LocalWikiEmbeddingConfig(
|
|
provider="local",
|
|
model="another-local-model",
|
|
local_model_path=str(tmp_path / "configured"),
|
|
dimensions=384,
|
|
)
|
|
effective = llmwiki_index_runtime.resolve_effective_local_wiki_index_config(
|
|
_source_config(configured),
|
|
LocalWikiEncoderSettings(
|
|
mode="local",
|
|
model_path=str(tmp_path / "uploaded"),
|
|
model_revision="uploaded-revision",
|
|
),
|
|
)
|
|
|
|
assert effective.embedding.provider == "local"
|
|
assert effective.embedding.model == "bge-embedding-m3"
|
|
assert effective.embedding.dimensions == 1024
|
|
assert effective.embedding.local_model_path == str(tmp_path / "uploaded")
|
|
|
|
|
|
def test_local_provider_never_uses_residual_base_url_as_readiness(tmp_path: Path) -> None:
|
|
config = LocalWikiIndexConfig(
|
|
enabled=True,
|
|
embedding=LocalWikiEmbeddingConfig(
|
|
provider="local",
|
|
model="bge-embedding-m3",
|
|
local_model_path=str(tmp_path / "missing"),
|
|
base_url="https://stale.example.test/v1",
|
|
dimensions=1024,
|
|
),
|
|
)
|
|
|
|
assert llmwiki_index_runtime._embedding_provider_ready(config) is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retired_embedding_rejects_future_work() -> None:
|
|
client = StrictWikiEmbeddingClient(
|
|
LocalWikiEmbeddingConfig(
|
|
provider="openai",
|
|
model="text-embedding-3-small",
|
|
base_url="https://api.openai.com/v1",
|
|
dimensions=1536,
|
|
)
|
|
)
|
|
|
|
await client.retire()
|
|
|
|
with pytest.raises(WikiEmbeddingError) as exc_info:
|
|
await client.embed_query("retry me")
|
|
assert getattr(exc_info.value, "category", None) == "runtime_retired"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_local_embedding_drains_worker_thread_before_releasing_lock(
|
|
monkeypatch,
|
|
) -> None:
|
|
client = StrictWikiEmbeddingClient(
|
|
LocalWikiEmbeddingConfig(
|
|
provider="local",
|
|
model="bge-embedding-m3",
|
|
dimensions=2,
|
|
)
|
|
)
|
|
started = Event()
|
|
release = Event()
|
|
|
|
def blocking_embed(_texts):
|
|
started.set()
|
|
release.wait(timeout=5)
|
|
return [[1.0, 0.0]]
|
|
|
|
monkeypatch.setattr(client, "_embed_local", blocking_embed)
|
|
embedding_task = asyncio.create_task(client.embed_texts(["text"]))
|
|
assert await asyncio.to_thread(started.wait, 1)
|
|
|
|
embedding_task.cancel()
|
|
await asyncio.sleep(0)
|
|
assert not embedding_task.done()
|
|
assert client._active_uses == 1
|
|
|
|
embedding_task.cancel()
|
|
await asyncio.sleep(0)
|
|
assert not embedding_task.done()
|
|
assert client._active_uses == 1
|
|
|
|
release.set()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await embedding_task
|
|
assert client._active_uses == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_status_is_available_for_remote_provider(monkeypatch) -> None:
|
|
configured = LocalWikiEmbeddingConfig(
|
|
provider="openai",
|
|
model="text-embedding-3-small",
|
|
base_url="https://api.openai.com/v1",
|
|
dimensions=1536,
|
|
)
|
|
app = FastAPI()
|
|
app.state.config = _source_config(configured)
|
|
app.state.llmwiki_effective_config = app.state.config.llmwiki.local_wiki_index
|
|
app.state.llmwiki_embedding = StrictWikiEmbeddingClient(configured)
|
|
request = Request({"type": "http", "headers": [], "app": app})
|
|
|
|
async def admin(_request):
|
|
return "admin", True
|
|
|
|
async def unchanged(_app, *, force=False):
|
|
return _app.state.llmwiki_effective_config
|
|
|
|
monkeypatch.setattr(llmwiki_index, "_actor", admin)
|
|
monkeypatch.setattr(llmwiki_index, "rebuild_local_wiki_index_runtime", unchanged)
|
|
monkeypatch.setattr(
|
|
llmwiki_index_runtime,
|
|
"load_system_settings",
|
|
lambda: SystemSettings(
|
|
local_wiki_encoder=LocalWikiEncoderSettings(mode="configured")
|
|
),
|
|
)
|
|
|
|
result = await llmwiki_index.local_encoder_status(request)
|
|
|
|
assert result["provider"] == "openai"
|
|
assert result["configured_provider"] == "openai"
|
|
assert result["mode"] == "configured"
|
|
assert result["using_local"] is False
|
|
assert result["state"] == "idle"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_local_path_save_and_restore_persist_mode_and_rebuild(
|
|
monkeypatch,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
model_dir = tmp_path / "bge-m3"
|
|
(model_dir / "onnx").mkdir(parents=True)
|
|
(model_dir / "onnx" / "model.onnx").write_bytes(b"onnx")
|
|
app = FastAPI()
|
|
request = Request({"type": "http", "headers": [], "app": app})
|
|
saved: list[SystemSettings] = []
|
|
current = SystemSettings()
|
|
rebuilds: list[bool] = []
|
|
|
|
async def admin(_request):
|
|
return "admin", True
|
|
|
|
async def rebuilt(_app, *, force=False):
|
|
rebuilds.append(force)
|
|
return LocalWikiIndexConfig()
|
|
|
|
monkeypatch.setattr(llmwiki_index, "_actor", admin)
|
|
monkeypatch.setattr(llmwiki_index, "load_system_settings", lambda: saved[-1] if saved else current)
|
|
monkeypatch.setattr(llmwiki_index, "save_system_settings", lambda value: saved.append(value) or value)
|
|
monkeypatch.setattr(llmwiki_index, "rebuild_local_wiki_index_runtime", rebuilt)
|
|
monkeypatch.setattr(
|
|
llmwiki_index,
|
|
"local_encoder_runtime_status",
|
|
lambda _app: {
|
|
"provider": "local" if saved[-1].local_wiki_encoder.mode == "local" else "openai",
|
|
"mode": saved[-1].local_wiki_encoder.mode,
|
|
},
|
|
)
|
|
|
|
local = await llmwiki_index.configure_local_encoder(
|
|
request,
|
|
llmwiki_index.LocalEncoderConfigRequest(model_path=str(model_dir)),
|
|
)
|
|
assert local["mode"] == "local"
|
|
assert saved[-1].local_wiki_encoder.mode == "local"
|
|
assert saved[-1].local_wiki_encoder.model_revision
|
|
|
|
configured = await llmwiki_index.use_configured_encoder(request)
|
|
assert configured["mode"] == "configured"
|
|
assert saved[-1].local_wiki_encoder.mode == "configured"
|
|
assert rebuilds == [True, True]
|
|
|
|
|
|
def test_both_shipped_offline_configs_enable_incremental_wiki_rescan() -> None:
|
|
offline_root = Path(__file__).resolve().parents[2]
|
|
for path in (offline_root / "config.yaml", offline_root / "backend" / "config.yaml"):
|
|
payload = yaml.safe_load(path.read_text(encoding="utf-8"))
|
|
local_index = payload["llmwiki"]["local_wiki_index"]
|
|
assert local_index["enabled"] is True
|
|
assert local_index["auto_sync"] is True
|