Initial curated RAGcore source from c16667f0c9397e2bfc8dd0255b114fba23c8d05e
This commit is contained in:
Public Source Publisher committed
commit
7ca726caf2
659 files changed
+98171
No files matched your search
@@ -0,0 +1,244 @@
|
||||
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
|
||||
Reference in new issue
Block a user