119 lines
4.0 KiB
Python
119 lines
4.0 KiB
Python
from __future__ import annotations
|
|
|
|
from uuid import uuid4
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from ragcore.application.models.bootstrap import (
|
|
EmbeddingBootstrapError,
|
|
EmbeddingBootstrapRequest,
|
|
ensure_embedding_profile,
|
|
)
|
|
from ragcore.application.models.registry import (
|
|
EmbeddingProfileNotFoundError,
|
|
EmbeddingProfileRegistry,
|
|
)
|
|
from ragcore.domain.control import DataPolicy
|
|
from ragcore.domain.models import EmbeddingProfile, EmbeddingProvider, NormalizationMode
|
|
from ragcore.infrastructure.models.memory import InMemoryEmbeddingProfileStore
|
|
|
|
|
|
class FakeClient:
|
|
def __init__(self, responses: list[httpx.Response]) -> None:
|
|
self.responses = responses
|
|
self.requests: list[str] = []
|
|
|
|
async def post(self, url: str, **kwargs: object) -> httpx.Response:
|
|
self.requests.append(url)
|
|
return self.responses.pop(0)
|
|
|
|
|
|
def _response(status: int, body: object) -> httpx.Response:
|
|
return httpx.Response(status, json=body, request=httpx.Request("POST", "http://test"))
|
|
|
|
|
|
def _request() -> EmbeddingBootstrapRequest:
|
|
return EmbeddingBootstrapRequest(
|
|
provider=EmbeddingProvider.OLLAMA,
|
|
model_identifier="nomic-embed-text:1.5",
|
|
base_url="http://ollama",
|
|
)
|
|
|
|
|
|
def _existing_profile() -> EmbeddingProfile:
|
|
return EmbeddingProfile(
|
|
id=uuid4(),
|
|
profile="hand-registered",
|
|
version="v1",
|
|
provider=EmbeddingProvider.OLLAMA,
|
|
model_identifier="mxbai-embed-large:1",
|
|
model_digest="a" * 64,
|
|
dimensions=3,
|
|
context_limit=4096,
|
|
tokenizer="sentencepiece-v1",
|
|
data_policy=DataPolicy.LOCAL_ONLY,
|
|
normalization=NormalizationMode.NONE,
|
|
)
|
|
|
|
|
|
async def test_noop_when_a_profile_is_already_active() -> None:
|
|
store = InMemoryEmbeddingProfileStore()
|
|
registry = EmbeddingProfileRegistry(store)
|
|
existing = _existing_profile()
|
|
await registry.register(existing)
|
|
await registry.activate(existing.id)
|
|
|
|
client = FakeClient([])
|
|
await ensure_embedding_profile(registry, _request(), http_client=client) # type: ignore[arg-type]
|
|
|
|
assert client.requests == []
|
|
assert (await registry.active()).id == existing.id
|
|
|
|
|
|
async def test_noop_and_no_raise_when_no_request_and_no_active_profile() -> None:
|
|
registry = EmbeddingProfileRegistry(InMemoryEmbeddingProfileStore())
|
|
client = FakeClient([])
|
|
|
|
await ensure_embedding_profile(registry, None, http_client=client) # type: ignore[arg-type]
|
|
|
|
with pytest.raises(EmbeddingProfileNotFoundError):
|
|
await registry.active()
|
|
|
|
|
|
async def test_bootstraps_a_real_profile_from_provider_and_model_probes() -> None:
|
|
registry = EmbeddingProfileRegistry(InMemoryEmbeddingProfileStore())
|
|
client = FakeClient(
|
|
[
|
|
_response(200, {"digest": "sha256:" + "b" * 64}),
|
|
_response(200, {"embeddings": [[0.1, 0.2, 0.3, 0.4]]}),
|
|
]
|
|
)
|
|
|
|
await ensure_embedding_profile(registry, _request(), http_client=client) # type: ignore[arg-type]
|
|
|
|
active = await registry.active()
|
|
assert active.model_identifier == "nomic-embed-text:1.5"
|
|
assert active.model_digest == "b" * 64
|
|
assert active.dimensions == 4
|
|
assert client.requests == ["http://ollama/api/show", "http://ollama/api/embed"]
|
|
|
|
|
|
async def test_bootstrap_raises_when_the_model_is_unreachable() -> None:
|
|
registry = EmbeddingProfileRegistry(InMemoryEmbeddingProfileStore())
|
|
client = FakeClient([_response(500, {"error": "down"})])
|
|
|
|
with pytest.raises(EmbeddingBootstrapError):
|
|
await ensure_embedding_profile(registry, _request(), http_client=client) # type: ignore[arg-type]
|
|
|
|
with pytest.raises(EmbeddingProfileNotFoundError):
|
|
await registry.active()
|
|
|
|
|
|
async def test_bootstrap_raises_on_malformed_digest() -> None:
|
|
registry = EmbeddingProfileRegistry(InMemoryEmbeddingProfileStore())
|
|
client = FakeClient([_response(200, {"digest": "not-a-digest"})])
|
|
|
|
with pytest.raises(EmbeddingBootstrapError):
|
|
await ensure_embedding_profile(registry, _request(), http_client=client) # type: ignore[arg-type]
|