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