mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-10-19 07:34:58 +00:00
84c7ffad18
The embedding model is now chosen in the last onboarding step and fixed for the whole library. Skipping keeps bge-small, which ships with the app and can't be deleted. Changing the model later is not built. - An embedding_indexes row records the model's spec and its vector table, and each document records which index embedded it. Migration 0021 stamps existing libraries with bge-small. Routes that queue embedding work return 409 embedding_not_chosen until a model is chosen. - Ingest, Studio and search embed through the active index; the old worker embedding module is removed. Workers wait for the schema, and documents interrupted when the app closed are marked failed. - Curated embedders (bge-small, granite-97m, granite-311m) live in the local manifest, run by a new onnxruntime engine. EMBEDDING is a catalog type, never a selection slot. - Hugging Face search finds ONNX embedders. Each download is checked: a probe sets its width, and a ten-question search check must pass. Small tokenizer.json files that Hugging Face lists without a checksum are hashed at the pinned commit. Results use the GGUF search's shapes. - Onboarding renders the embedding step with the same component and shared list as the other steps; only its choice hook differs. Use marks the choice and Finish sends it. Bundled models have no Delete. - Settings has a read-only Embedding section. - ADRs 0036 and 0037, a new embedding architecture doc, and updates to the docs it touches. New strings are translated into every language.
306 lines
11 KiB
Python
306 lines
11 KiB
Python
import pytest
|
|
from httpx import AsyncClient
|
|
from sqlalchemy import Engine
|
|
|
|
from modules.documents.models import Document, DocumentStatus, DocumentType
|
|
from modules.embedding.active import active_index
|
|
from modules.embedding.bundled import BGE
|
|
from modules.llm import providers as registry
|
|
from modules.llm.activity import model_activity, model_key
|
|
from modules.llm.catalog.local.dependencies import get_local_catalog
|
|
from modules.llm.providers.types import Message, Model
|
|
from modules.workspaces.models import Workspace
|
|
from shared.db import create_session_factory
|
|
|
|
pytestmark = pytest.mark.integration
|
|
|
|
|
|
async def test_providers_report_health_and_capability(
|
|
client: AsyncClient, llamacpp_server: str
|
|
) -> None:
|
|
"""The download UI is shown by capability, not by a provider's name."""
|
|
body = (await client.get("/llm/providers")).json()
|
|
|
|
local = next(entry for entry in body if entry["name"] == "llamacpp")
|
|
assert local["healthy"] is True
|
|
assert local["can_download"] is True
|
|
|
|
|
|
async def test_installed_models_carry_their_capabilities(
|
|
client: AsyncClient, llamacpp_server: str
|
|
) -> None:
|
|
"""A chosen model has to be one the runtime actually holds."""
|
|
body = (await client.get("/llm/providers/llamacpp/models")).json()
|
|
|
|
assert {model["name"] for model in body} == {"Qwen3-1.7B-Q4_K_M", "Qwen3-4B-Q4_K_M"}
|
|
assert "text_gen" in body[0]["selectable_for"]
|
|
|
|
|
|
async def test_an_unknown_provider_is_a_404(client: AsyncClient) -> None:
|
|
"""A path names a provider the registry does not have."""
|
|
assert (await client.get("/llm/providers/openai/models")).status_code == 404
|
|
|
|
|
|
async def test_the_selection_is_read_after_it_is_set(
|
|
client: AsyncClient, llamacpp_server: str
|
|
) -> None:
|
|
"""Set persists the choice; read is how the rest of the app learns it."""
|
|
assert (await client.get("/llm/selection/text_gen")).status_code == 404
|
|
|
|
written = await client.put(
|
|
"/llm/selection/text_gen",
|
|
json={"provider": "llamacpp", "name": "Qwen3-4B-Q4_K_M"},
|
|
)
|
|
assert written.status_code == 200
|
|
assert written.json()["name"] == "Qwen3-4B-Q4_K_M"
|
|
assert written.json()["model_type"] == "text_gen"
|
|
|
|
read = await client.get("/llm/selection/text_gen")
|
|
assert read.json() == written.json()
|
|
|
|
|
|
async def test_the_selection_says_which_prompt_tier_the_model_gets(
|
|
client: AsyncClient, llamacpp_server: str
|
|
) -> None:
|
|
"""A 1.7B model asks for a different prompt than a hosted frontier one."""
|
|
await client.put(
|
|
"/llm/selection/text_gen",
|
|
json={"provider": "llamacpp", "name": "Qwen3-1.7B-Q4_K_M"},
|
|
)
|
|
|
|
read = (await client.get("/llm/selection/text_gen")).json()
|
|
|
|
assert read["tier"] == "compact"
|
|
|
|
|
|
async def test_selecting_a_chat_model_does_not_complete_onboarding(
|
|
client: AsyncClient, llamacpp_server: str
|
|
) -> None:
|
|
"""Use persists the chat model; Continue writes the completion marker."""
|
|
assert (await client.get("/llm/onboarding")).json() == {"completed": False}
|
|
|
|
await client.put(
|
|
"/llm/selection/text_gen",
|
|
json={"provider": "llamacpp", "name": "Qwen3-1.7B-Q4_K_M"},
|
|
)
|
|
assert (await client.get("/llm/onboarding")).json() == {"completed": False}
|
|
|
|
completed = await client.post("/llm/onboarding")
|
|
assert completed.status_code == 200
|
|
assert completed.json() == {"completed": True}
|
|
|
|
|
|
async def test_the_embedder_cannot_be_selected_like_a_model(
|
|
client: AsyncClient,
|
|
) -> None:
|
|
"""It is fixed with the index at onboarding; a selection would swap it."""
|
|
reply = await client.put(
|
|
"/llm/selection/embedding",
|
|
json={
|
|
"provider": "onnxruntime",
|
|
"name": "granite-embedding-97m-multilingual-r2",
|
|
},
|
|
)
|
|
|
|
assert reply.status_code == 422
|
|
assert (await client.get("/llm/selection/embedding")).status_code == 422
|
|
|
|
|
|
async def test_finishing_onboarding_fixes_bge_as_the_embedder(
|
|
unlocked_client: AsyncClient, unlocked_engine: Engine, llamacpp_server: str
|
|
) -> None:
|
|
"""Nothing was chosen, so the bundled model is; finishing twice changes nothing."""
|
|
client = unlocked_client
|
|
await client.put(
|
|
"/llm/selection/text_gen",
|
|
json={"provider": "llamacpp", "name": "Qwen3-1.7B-Q4_K_M"},
|
|
)
|
|
|
|
await client.post("/llm/onboarding")
|
|
with create_session_factory(unlocked_engine)() as session:
|
|
first = active_index(session)
|
|
assert (await client.post("/llm/onboarding")).status_code == 200
|
|
with create_session_factory(unlocked_engine)() as session:
|
|
again = active_index(session)
|
|
|
|
assert first is not None
|
|
assert first.spec == BGE
|
|
assert again == first
|
|
|
|
|
|
async def test_onboarding_cannot_complete_without_a_chat_model(
|
|
client: AsyncClient,
|
|
) -> None:
|
|
"""Onboarding needs a chat model chosen before it can finish."""
|
|
reply = await client.post("/llm/onboarding")
|
|
assert reply.status_code == 422
|
|
assert (await client.get("/llm/onboarding")).json() == {"completed": False}
|
|
|
|
|
|
async def test_deleting_the_selected_local_model_clears_only_the_selection(
|
|
client: AsyncClient, llamacpp_server: str
|
|
) -> None:
|
|
"""Deleting the active model preserves completed onboarding."""
|
|
await client.put(
|
|
"/llm/selection/text_gen",
|
|
json={"provider": "llamacpp", "name": "Qwen3-1.7B-Q4_K_M"},
|
|
)
|
|
assert (await client.post("/llm/onboarding")).status_code == 200
|
|
|
|
deleted = await client.delete("/llm/models/Qwen3-1.7B-Q4_K_M")
|
|
|
|
assert deleted.status_code == 200
|
|
assert deleted.json() == {
|
|
"name": "Qwen3-1.7B-Q4_K_M",
|
|
"selection_cleared": True,
|
|
}
|
|
assert (await client.get("/llm/selection/text_gen")).status_code == 404
|
|
assert (await client.get("/llm/onboarding")).json() == {"completed": True}
|
|
# Disk is the inventory, and it is what the catalog reads. The router still
|
|
# lists the model until its next restart, which the preset rewrite triggers:
|
|
# it scans its directory once at startup and has no way to be told otherwise.
|
|
rows = (await client.get("/llm/catalog/local")).json()["rows"]
|
|
installed = [
|
|
b["installed_as"] for row in rows for b in row["builds"] if b["installed_as"]
|
|
]
|
|
assert installed == ["Qwen3-4B-Q4_K_M"]
|
|
|
|
|
|
async def test_remote_models_cannot_be_deleted(client: AsyncClient) -> None:
|
|
"""Remote connections never enter local provider storage routes."""
|
|
reply = await client.delete(
|
|
"/llm/providers/openai_compatible/models/anthropic%2Fclaude-3.5-sonnet"
|
|
)
|
|
|
|
assert reply.status_code == 404
|
|
|
|
|
|
async def test_a_model_in_use_cannot_be_deleted(
|
|
client: AsyncClient, llamacpp_server: str
|
|
) -> None:
|
|
"""Deletion cannot race an active generation."""
|
|
key = model_key("llamacpp", "Qwen3-1.7B-Q4_K_M")
|
|
await model_activity.acquire_use(key)
|
|
try:
|
|
reply = await client.delete("/llm/models/Qwen3-1.7B-Q4_K_M")
|
|
finally:
|
|
await model_activity.release_use(key)
|
|
|
|
assert reply.status_code == 409
|
|
assert reply.json()["detail"] == "model is currently in use"
|
|
|
|
|
|
async def test_a_model_cannot_be_deleted_while_one_is_installing(
|
|
client: AsyncClient, llamacpp_server: str
|
|
) -> None:
|
|
"""One runtime mutation cannot overlap another."""
|
|
lock = get_local_catalog().install_lock()
|
|
await lock.acquire()
|
|
try:
|
|
reply = await client.delete("/llm/models/Qwen3-1.7B-Q4_K_M")
|
|
finally:
|
|
lock.release()
|
|
|
|
assert reply.status_code == 409
|
|
assert reply.json()["detail"] == "a model is already being installed"
|
|
|
|
|
|
async def test_a_model_cannot_be_deleted_while_studio_is_generating(
|
|
client: AsyncClient, engine: Engine, llamacpp_server: str
|
|
) -> None:
|
|
"""The API sees model work running in the separate Studio worker."""
|
|
with create_session_factory(engine)() as session:
|
|
workspace = Workspace(name="Studio")
|
|
session.add(workspace)
|
|
session.flush()
|
|
session.add(
|
|
Document(
|
|
workspace_id=workspace.id,
|
|
title="Running artifact",
|
|
document_type=DocumentType.ARTIFACT,
|
|
status=DocumentStatus.PROCESSING,
|
|
error_message=None,
|
|
content=None,
|
|
content_hash=None,
|
|
dedup_key=None,
|
|
document_metadata=None,
|
|
)
|
|
)
|
|
session.commit()
|
|
|
|
reply = await client.delete("/llm/models/Qwen3-1.7B-Q4_K_M")
|
|
|
|
assert reply.status_code == 409
|
|
assert (
|
|
reply.json()["detail"] == "a model cannot be deleted while Studio is generating"
|
|
)
|
|
|
|
|
|
async def test_choosing_again_updates_in_place(
|
|
client: AsyncClient, llamacpp_server: str
|
|
) -> None:
|
|
"""One row per role: the second choice replaces the first, not adds to it."""
|
|
await client.put(
|
|
"/llm/selection/text_gen",
|
|
json={"provider": "llamacpp", "name": "Qwen3-1.7B-Q4_K_M"},
|
|
)
|
|
await client.put(
|
|
"/llm/selection/text_gen",
|
|
json={"provider": "llamacpp", "name": "Qwen3-4B-Q4_K_M"},
|
|
)
|
|
|
|
assert (await client.get("/llm/selection/text_gen")).json()[
|
|
"name"
|
|
] == "Qwen3-4B-Q4_K_M"
|
|
|
|
|
|
async def test_a_selection_names_a_known_provider(client: AsyncClient) -> None:
|
|
"""A choice pointing at no provider would never resolve to a model."""
|
|
reply = await client.put(
|
|
"/llm/selection/text_gen", json={"provider": "openai", "name": "gpt-4o"}
|
|
)
|
|
assert reply.status_code == 422
|
|
|
|
|
|
async def test_a_selection_names_an_installed_model(
|
|
client: AsyncClient, llamacpp_server: str
|
|
) -> None:
|
|
"""A stale or invented model name is rejected before it reaches chat."""
|
|
reply = await client.put(
|
|
"/llm/selection/text_gen",
|
|
json={"provider": "llamacpp", "name": "does-not-exist"},
|
|
)
|
|
|
|
assert reply.status_code == 422
|
|
assert reply.json()["detail"] == "model is not installed: does-not-exist"
|
|
assert (await client.get("/llm/selection/text_gen")).status_code == 404
|
|
|
|
|
|
async def test_a_text_selection_requires_a_text_model(
|
|
client: AsyncClient, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
"""An installed embedding model cannot be selected to answer chat."""
|
|
|
|
class EmbeddingOnly:
|
|
name = "llamacpp"
|
|
|
|
async def health(self) -> bool:
|
|
return True
|
|
|
|
async def models(self) -> list[Model]:
|
|
# Known, and a type no slot holds: an embedder.
|
|
return [Model("embedder", installed=True, types=(), known=True)]
|
|
|
|
def chat(self, model: str, messages: list[Message]): # pragma: no cover
|
|
raise NotImplementedError
|
|
|
|
monkeypatch.setitem(registry.REGISTRY, "llamacpp", EmbeddingOnly)
|
|
|
|
reply = await client.put(
|
|
"/llm/selection/text_gen",
|
|
json={"provider": "llamacpp", "name": "embedder"},
|
|
)
|
|
|
|
assert reply.status_code == 422
|
|
assert reply.json()["detail"] == "model does not support text_gen: embedder"
|
|
assert (await client.get("/llm/selection/text_gen")).status_code == 404
|