245 lines
8.4 KiB
Python
245 lines
8.4 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import math
|
|
from dataclasses import replace
|
|
from uuid import uuid4
|
|
|
|
import httpx
|
|
import pytest
|
|
from opentelemetry.sdk.metrics import MeterProvider
|
|
from opentelemetry.sdk.metrics.export import InMemoryMetricReader
|
|
from opentelemetry.sdk.resources import Resource
|
|
|
|
from ragcore.application.limits import ConcurrencyBudget
|
|
from ragcore.application.models import EmbeddingProfileRegistry
|
|
from ragcore.application.models.registry import EmbeddingProfileNotFoundError
|
|
from ragcore.domain.control import DataPolicy
|
|
from ragcore.domain.models import EmbeddingProfile, EmbeddingProvider, NormalizationMode
|
|
from ragcore.infrastructure.models import InMemoryEmbeddingProfileStore, OllamaEmbeddingAdapter
|
|
from ragcore.infrastructure.models.ollama import EmbeddingProviderError
|
|
from ragcore.infrastructure.models.persistence import JsonEmbeddingProfileStore
|
|
from ragcore.platform.telemetry import TelemetryMetrics
|
|
|
|
|
|
def profile(**changes: object) -> EmbeddingProfile:
|
|
value: dict[str, object] = {
|
|
"id": uuid4(),
|
|
"profile": "dense",
|
|
"version": "v1",
|
|
"provider": EmbeddingProvider.OLLAMA,
|
|
"model_identifier": "nomic-embed-text:1.5",
|
|
"model_digest": "a" * 64,
|
|
"dimensions": 2,
|
|
"context_limit": 4096,
|
|
"tokenizer": "sentencepiece-v1",
|
|
"data_policy": DataPolicy.LOCAL_ONLY,
|
|
"max_batch_size": 2,
|
|
"max_retries": 2,
|
|
"normalization": NormalizationMode.L2,
|
|
}
|
|
value.update(changes)
|
|
return EmbeddingProfile(**value)
|
|
|
|
|
|
class FakeClient:
|
|
def __init__(self, responses: list[httpx.Response]) -> None:
|
|
self.responses = responses
|
|
self.requests: list[dict[str, object]] = []
|
|
|
|
async def post(self, url: str, **kwargs: object) -> httpx.Response:
|
|
self.requests.append({"url": url, **kwargs})
|
|
return self.responses.pop(0)
|
|
|
|
|
|
async def no_sleep(_: float) -> None:
|
|
return None
|
|
|
|
|
|
def response(status: int, body: object) -> httpx.Response:
|
|
return httpx.Response(status, json=body, request=httpx.Request("POST", "http://test"))
|
|
|
|
|
|
|
|
def test_adapter_rejects_nonlocal_provider_hosts() -> None:
|
|
with pytest.raises(ValueError, match="allowlisted local"):
|
|
OllamaEmbeddingAdapter(
|
|
profile(),
|
|
FakeClient([]), # type: ignore[arg-type]
|
|
base_url="http://169.254.169.254",
|
|
)
|
|
def test_profile_is_immutable_and_rejects_floating_identifiers() -> None:
|
|
value = profile()
|
|
with pytest.raises(AttributeError):
|
|
value.model_identifier = "nomic-embed-text:latest" # type: ignore[misc]
|
|
with pytest.raises(ValueError, match="floating"):
|
|
profile(model_identifier="nomic-embed-text:latest")
|
|
with pytest.raises(ValueError, match="digest"):
|
|
profile(model_digest="not-a-digest")
|
|
with pytest.raises(ValueError, match="dimensions"):
|
|
profile(dimensions=0)
|
|
|
|
|
|
|
|
@pytest.mark.integration
|
|
async def test_json_store_persists_profile_and_exact_active_id(tmp_path) -> None:
|
|
path = tmp_path / "models.json"
|
|
first = profile()
|
|
store = JsonEmbeddingProfileStore(path)
|
|
await store.add(first)
|
|
await store.activate(first.id)
|
|
|
|
recovered = JsonEmbeddingProfileStore(path)
|
|
assert await recovered.get(first.id) == first
|
|
assert await recovered.get_active() == first
|
|
|
|
@pytest.mark.integration
|
|
async def test_registry_activates_by_exact_profile_id_only() -> None:
|
|
store = InMemoryEmbeddingProfileStore()
|
|
registry = EmbeddingProfileRegistry(store)
|
|
first = profile()
|
|
second = profile(version="v2")
|
|
await registry.register(first)
|
|
with pytest.raises(ValueError, match="logical revision"):
|
|
await registry.register(replace(first, id=uuid4(), model_digest="b" * 64))
|
|
await registry.register(second)
|
|
await registry.activate(second.id)
|
|
assert (await registry.active()).id == second.id
|
|
with pytest.raises(EmbeddingProfileNotFoundError):
|
|
await registry.activate(uuid4())
|
|
|
|
|
|
@pytest.mark.integration
|
|
async def test_adapter_batches_prefixes_and_normalizes() -> None:
|
|
client = FakeClient(
|
|
[
|
|
response(200, {"embeddings": [[3.0, 4.0], [0.0, 2.0]]}),
|
|
response(200, {"embeddings": [[1.0, 0.0]]}),
|
|
]
|
|
)
|
|
adapter = OllamaEmbeddingAdapter(
|
|
profile(document_prefix="passage: "),
|
|
client, # type: ignore[arg-type]
|
|
base_url="http://ollama",
|
|
sleep=no_sleep,
|
|
)
|
|
vectors = await adapter.embed(("one", "two", "three"))
|
|
assert len(vectors) == 3
|
|
assert all(math.isclose(math.sqrt(sum(x * x for x in vector)), 1.0) for vector in vectors)
|
|
assert len(client.requests) == 2
|
|
assert client.requests[0]["json"] == {
|
|
"model": "nomic-embed-text:1.5",
|
|
"input": ["passage: one", "passage: two"],
|
|
}
|
|
|
|
|
|
@pytest.mark.integration
|
|
async def test_adapter_bounds_retries_timeout_and_response_shape() -> None:
|
|
client = FakeClient(
|
|
[
|
|
response(503, {"error": "busy"}),
|
|
response(429, {"error": "busy"}),
|
|
response(200, {"embeddings": [[1.0, 0.0]]}),
|
|
]
|
|
)
|
|
adapter = OllamaEmbeddingAdapter(profile(), client, base_url="http://ollama", sleep=no_sleep) # type: ignore[arg-type]
|
|
assert await adapter.embed(("one",)) == ((1.0, 0.0),)
|
|
assert len(client.requests) == 3
|
|
|
|
malformed = OllamaEmbeddingAdapter(
|
|
profile(),
|
|
FakeClient([response(200, {"embeddings": [[1.0]]})]), # type: ignore[arg-type]
|
|
base_url="http://ollama",
|
|
sleep=no_sleep,
|
|
)
|
|
with pytest.raises(EmbeddingProviderError, match="dimensions"):
|
|
await malformed.embed(("one",))
|
|
|
|
cardinality = OllamaEmbeddingAdapter(
|
|
profile(),
|
|
FakeClient([response(200, {"embeddings": [[1.0, 0.0]]})]), # type: ignore[arg-type]
|
|
base_url="http://ollama",
|
|
sleep=no_sleep,
|
|
)
|
|
with pytest.raises(EmbeddingProviderError, match="cardinality"):
|
|
await cardinality.embed(("one", "two"))
|
|
|
|
|
|
@pytest.mark.integration
|
|
async def test_adapter_surfaces_concurrency_exhaustion_as_a_provider_error() -> None:
|
|
budget = ConcurrencyBudget(name="embedding-calls", max_concurrent=1, max_queued=0)
|
|
release = asyncio.Event()
|
|
|
|
async def occupy() -> None:
|
|
async with budget.acquire():
|
|
await release.wait()
|
|
|
|
holder = asyncio.create_task(occupy())
|
|
await asyncio.sleep(0)
|
|
|
|
adapter = OllamaEmbeddingAdapter(
|
|
profile(),
|
|
FakeClient([]), # type: ignore[arg-type]
|
|
base_url="http://ollama",
|
|
concurrency_budget=budget,
|
|
)
|
|
with pytest.raises(EmbeddingProviderError, match="concurrency"):
|
|
await adapter.embed(("one",))
|
|
|
|
release.set()
|
|
await holder
|
|
|
|
|
|
@pytest.mark.integration
|
|
async def test_adapter_rejects_zero_vector_and_unbounded_input() -> None:
|
|
zero = OllamaEmbeddingAdapter(
|
|
profile(),
|
|
FakeClient([response(200, {"embeddings": [[0.0, 0.0]]})]), # type: ignore[arg-type]
|
|
base_url="http://ollama",
|
|
sleep=no_sleep,
|
|
)
|
|
with pytest.raises(EmbeddingProviderError, match="zero"):
|
|
await zero.embed(("one",))
|
|
with pytest.raises(ValueError, match="input"):
|
|
await zero.embed(("x",) * 4_097)
|
|
await zero.embed(("x" * 32_001,))
|
|
|
|
|
|
@pytest.mark.integration
|
|
async def test_adapter_records_provider_duration_on_success_and_error_on_failure() -> None:
|
|
reader = InMemoryMetricReader()
|
|
provider = MeterProvider(
|
|
resource=Resource.create({"service.name": "test"}), metric_readers=[reader]
|
|
)
|
|
metrics = TelemetryMetrics(meter_provider=provider)
|
|
|
|
ok_adapter = OllamaEmbeddingAdapter(
|
|
profile(),
|
|
FakeClient([response(200, {"embeddings": [[1.0, 0.0]]})]), # type: ignore[arg-type]
|
|
base_url="http://ollama",
|
|
sleep=no_sleep,
|
|
metrics=metrics,
|
|
)
|
|
await ok_adapter.embed(("one",))
|
|
|
|
failing_adapter = OllamaEmbeddingAdapter(
|
|
profile(),
|
|
FakeClient([response(200, {"embeddings": [[1.0]]})]), # type: ignore[arg-type]
|
|
base_url="http://ollama",
|
|
sleep=no_sleep,
|
|
metrics=metrics,
|
|
)
|
|
with pytest.raises(EmbeddingProviderError):
|
|
await failing_adapter.embed(("one",))
|
|
|
|
data = reader.get_metrics_data()
|
|
assert data is not None
|
|
names = {
|
|
metric.name
|
|
for resource_metric in data.resource_metrics
|
|
for scope_metric in resource_metric.scope_metrics
|
|
for metric in scope_metric.metrics
|
|
}
|
|
assert "ragcore_provider_duration_seconds" in names
|
|
assert "ragcore_provider_errors_total" in names
|