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
|
||||
@@ -0,0 +1,225 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from ragcore.config import Settings
|
||||
from ragcore.domain.control import DataPolicy
|
||||
from ragcore.domain.models import EmbeddingProfile, EmbeddingProvider
|
||||
from ragcore.infrastructure.models.embedding_factory import (
|
||||
EmbeddingAdapterFactory,
|
||||
create_embedding_adapter_factory,
|
||||
)
|
||||
from ragcore.infrastructure.models.ollama import EmbeddingProviderError
|
||||
|
||||
|
||||
def ollama_profile() -> EmbeddingProfile:
|
||||
return EmbeddingProfile(
|
||||
id=uuid4(),
|
||||
profile="test-ollama",
|
||||
version="v1",
|
||||
provider=EmbeddingProvider.OLLAMA,
|
||||
model_identifier="nomic-embed-text:v1",
|
||||
model_digest="a" * 64,
|
||||
dimensions=2,
|
||||
context_limit=8192,
|
||||
tokenizer="sentencepiece-v1",
|
||||
data_policy=DataPolicy.LOCAL_ONLY,
|
||||
document_prefix="doc: ",
|
||||
query_prefix="query: ",
|
||||
)
|
||||
|
||||
|
||||
def modelforge_profile(**changes: object) -> EmbeddingProfile:
|
||||
values: dict[str, object] = {
|
||||
"id": uuid4(),
|
||||
"profile": "test-modelforge",
|
||||
"version": "v1",
|
||||
"provider": EmbeddingProvider.MODELFORGE,
|
||||
"model_identifier": "rag.embedding@1",
|
||||
"model_digest": "b" * 64,
|
||||
"dimensions": 3,
|
||||
"context_limit": 8192,
|
||||
"tokenizer": "capability-managed",
|
||||
"data_policy": DataPolicy.LOCAL_ONLY,
|
||||
"max_retries": 0,
|
||||
"embedding_space_id": uuid4(),
|
||||
}
|
||||
values.update(changes)
|
||||
return EmbeddingProfile(**values)
|
||||
|
||||
|
||||
def gateway_response(profile: EmbeddingProfile, count: int) -> dict[str, object]:
|
||||
return {
|
||||
"capability": "rag.embedding@1",
|
||||
"dimension": profile.dimensions,
|
||||
"normalized": True,
|
||||
"embedding_space_id": str(profile.embedding_space_id),
|
||||
"data": [[1.0, 0.0, 0.0] for _ in range(count)],
|
||||
"request_id": str(uuid4()),
|
||||
}
|
||||
|
||||
|
||||
def factory(
|
||||
client: httpx.AsyncClient,
|
||||
*,
|
||||
credential_file: Path,
|
||||
modelforge_base_url: str = "http://modelforge:8000",
|
||||
experiment_route: str | None = None,
|
||||
) -> EmbeddingAdapterFactory:
|
||||
settings = Settings(
|
||||
env="test",
|
||||
ollama_url="http://ollama:11434",
|
||||
modelforge_base_url=modelforge_base_url,
|
||||
modelforge_credential_file=credential_file,
|
||||
modelforge_embedding_capability="rag.embedding@1",
|
||||
modelforge_embedding_experiment_route=experiment_route,
|
||||
)
|
||||
return create_embedding_adapter_factory(settings=settings, http_client=client)
|
||||
|
||||
|
||||
def write_credential(path: Path) -> None:
|
||||
path.write_text("mf_" + "x" * 48, encoding="utf-8")
|
||||
path.chmod(0o600)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_factory_routes_modelforge_profile_to_pinned_capability(tmp_path: Path) -> None:
|
||||
profile = modelforge_profile()
|
||||
credential_file = tmp_path / "credential"
|
||||
write_credential(credential_file)
|
||||
seen: list[httpx.Request] = []
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
seen.append(request)
|
||||
body = json.loads(request.content)
|
||||
return httpx.Response(200, json=gateway_response(profile, len(body["input"])))
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.MockTransport(respond), follow_redirects=False
|
||||
) as client:
|
||||
adapter = factory(client, credential_file=credential_file)(profile)
|
||||
vectors = await adapter.embed(["question"], query=True)
|
||||
|
||||
assert vectors == ((1.0, 0.0, 0.0),)
|
||||
assert len(seen) == 1
|
||||
assert seen[0].url.host == "modelforge"
|
||||
assert seen[0].url.path == "/api/v1/capabilities/rag.embedding@1/invoke"
|
||||
assert json.loads(seen[0].content) == {"input": ["question"], "input_type": "query"}
|
||||
assert seen[0].headers["authorization"].startswith("Bearer mf_")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_factory_routes_exact_selection_experiment_key(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
profile = modelforge_profile()
|
||||
credential_file = tmp_path / "credential"
|
||||
write_credential(credential_file)
|
||||
seen: list[httpx.Request] = []
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
seen.append(request)
|
||||
body = json.loads(request.content)
|
||||
return httpx.Response(200, json=gateway_response(profile, len(body["input"])))
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.MockTransport(respond), follow_redirects=False
|
||||
) as client:
|
||||
adapter = factory(
|
||||
client,
|
||||
credential_file=credential_file,
|
||||
experiment_route="ignored-global-route",
|
||||
)(profile, experiment_route_key="rc-canary-2026-09")
|
||||
vectors = await adapter.embed(["document"])
|
||||
|
||||
assert vectors == ((1.0, 0.0, 0.0),)
|
||||
assert len(seen) == 1
|
||||
assert seen[0].url.path == ("/api/v1/capability-experiments/rc-canary-2026-09/invoke")
|
||||
assert json.loads(seen[0].content) == {"input": ["document"], "input_type": "document"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_factory_rejects_an_unsafe_configured_experiment_route(tmp_path: Path) -> None:
|
||||
def unused_handler(_request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(500)
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.MockTransport(unused_handler), follow_redirects=False
|
||||
) as client:
|
||||
configured = factory(
|
||||
client,
|
||||
credential_file=tmp_path / "credential",
|
||||
experiment_route="valid-global-route",
|
||||
)
|
||||
with pytest.raises(ValueError, match="experiment_route"):
|
||||
configured(modelforge_profile(), experiment_route_key="../production")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_modelforge_credential_fails_before_any_provider_request(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
seen: list[httpx.Request] = []
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
seen.append(request)
|
||||
return httpx.Response(200, json={"embeddings": [[0.0, 1.0]]})
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.MockTransport(respond), follow_redirects=False
|
||||
) as client:
|
||||
adapter = factory(client, credential_file=tmp_path / "missing")(modelforge_profile())
|
||||
with pytest.raises(EmbeddingProviderError, match="credential file"):
|
||||
await adapter.embed(["question"])
|
||||
|
||||
assert seen == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_modelforge_gateway_never_falls_back_to_ollama(tmp_path: Path) -> None:
|
||||
profile = modelforge_profile()
|
||||
credential_file = tmp_path / "credential"
|
||||
write_credential(credential_file)
|
||||
seen: list[str] = []
|
||||
|
||||
def unavailable(request: httpx.Request) -> httpx.Response:
|
||||
seen.append(request.url.host + request.url.path)
|
||||
return httpx.Response(503)
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.MockTransport(unavailable), follow_redirects=False
|
||||
) as client:
|
||||
adapter = factory(client, credential_file=credential_file)(profile)
|
||||
with pytest.raises(EmbeddingProviderError, match="unavailable"):
|
||||
await adapter.embed(["question"])
|
||||
|
||||
assert seen == ["modelforge/api/v1/capabilities/rag.embedding@1/invoke"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_factory_keeps_ollama_request_model_and_prefix_contract(tmp_path: Path) -> None:
|
||||
seen: list[httpx.Request] = []
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
seen.append(request)
|
||||
return httpx.Response(200, json={"embeddings": [[0.6, 0.8]]})
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.MockTransport(respond), follow_redirects=False
|
||||
) as client:
|
||||
adapter = factory(client, credential_file=tmp_path / "unused")(ollama_profile())
|
||||
vectors = await adapter.embed(["question"], query=True)
|
||||
|
||||
assert vectors == ((0.6, 0.8),)
|
||||
assert len(seen) == 1
|
||||
assert seen[0].url.host == "ollama"
|
||||
assert seen[0].url.path == "/api/embed"
|
||||
assert json.loads(seen[0].content) == {
|
||||
"model": "nomic-embed-text:v1",
|
||||
"input": ["query: question"],
|
||||
}
|
||||
@@ -0,0 +1,312 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from ragcore.domain.control import DataPolicy
|
||||
from ragcore.domain.models import EmbeddingProfile, EmbeddingProvider
|
||||
from ragcore.infrastructure.models import EmbeddingProviderError, ModelForgeEmbeddingAdapter
|
||||
|
||||
|
||||
def profile() -> EmbeddingProfile:
|
||||
return EmbeddingProfile(
|
||||
id=uuid4(),
|
||||
profile="modelforge-rag",
|
||||
version="v1",
|
||||
provider=EmbeddingProvider.MODELFORGE,
|
||||
model_identifier="rag.embedding@1",
|
||||
model_digest="a" * 64,
|
||||
dimensions=3,
|
||||
context_limit=8192,
|
||||
tokenizer="capability-managed",
|
||||
data_policy=DataPolicy.LOCAL_ONLY,
|
||||
max_batch_size=2,
|
||||
max_retries=1,
|
||||
embedding_space_id=uuid4(),
|
||||
)
|
||||
|
||||
|
||||
def credential(tmp_path: Path) -> Path:
|
||||
path = tmp_path / "credential"
|
||||
path.write_text("mf_" + "x" * 48, encoding="utf-8")
|
||||
path.chmod(0o600)
|
||||
return path
|
||||
|
||||
|
||||
def response(item: EmbeddingProfile, vectors: list[list[float]]) -> dict[str, object]:
|
||||
return {
|
||||
"capability": "rag.embedding@1",
|
||||
"dimension": item.dimensions,
|
||||
"normalized": True,
|
||||
"embedding_space_id": str(item.embedding_space_id),
|
||||
"data": vectors,
|
||||
"request_id": str(uuid4()),
|
||||
"execution": {
|
||||
"cold": False,
|
||||
"node": "Tower",
|
||||
"residency": "warm",
|
||||
"load_count": 1,
|
||||
"timings": {
|
||||
"gateway_ms": 1,
|
||||
"queue_ms": 0,
|
||||
"load_ms": 0,
|
||||
"inference_ms": 2,
|
||||
"total_ms": 3,
|
||||
},
|
||||
},
|
||||
"usage": {"input_count": len(vectors), "input_tokens": 1},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_capability_first_batch_auth_space_and_correlation(tmp_path: Path) -> None:
|
||||
item = profile()
|
||||
seen: list[httpx.Request] = []
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
seen.append(request)
|
||||
count = len(json.loads(request.content)["input"])
|
||||
return httpx.Response(200, json=response(item, [[1.0, 0.0, 0.0]] * count))
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
||||
adapter = ModelForgeEmbeddingAdapter(
|
||||
item, client, base_url="http://modelforge:8000", credential_file=credential(tmp_path)
|
||||
)
|
||||
vectors, evidence = await adapter.embed_with_evidence(["one", "two", "three"])
|
||||
assert len(vectors) == 3
|
||||
assert len(seen) == 2
|
||||
assert all(
|
||||
request.url.path.endswith("/capabilities/rag.embedding@1/invoke") for request in seen
|
||||
)
|
||||
assert all(request.headers["authorization"].startswith("Bearer mf_") for request in seen)
|
||||
assert all(request.headers["x-correlation-id"].startswith("ragcore-") for request in seen)
|
||||
assert all(
|
||||
b"model" not in request.content and b"Tower" not in request.content for request in seen
|
||||
)
|
||||
assert evidence.total_ms >= evidence.gateway_round_trip_ms >= 0
|
||||
assert evidence.preparation_ms >= 0 and evidence.response_validation_ms >= 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("failure", ["capability", "space", "dimension", "nan"])
|
||||
async def test_response_contract_fails_closed(tmp_path: Path, failure: str) -> None:
|
||||
item = profile()
|
||||
body = response(item, [[1.0, 0.0, 0.0]])
|
||||
if failure == "capability":
|
||||
body["capability"] = "other.embedding@1"
|
||||
elif failure == "space":
|
||||
body["embedding_space_id"] = str(uuid4())
|
||||
elif failure == "dimension":
|
||||
body["dimension"] = 4
|
||||
else:
|
||||
body["data"] = [[float("nan"), 0.0, 0.0]]
|
||||
|
||||
def handler(_request: httpx.Request) -> httpx.Response:
|
||||
if failure == "nan":
|
||||
raw = json.dumps(body, allow_nan=True).encode()
|
||||
return httpx.Response(200, content=raw, headers={"content-type": "application/json"})
|
||||
return httpx.Response(200, json=body)
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
||||
adapter = ModelForgeEmbeddingAdapter(
|
||||
item, client, base_url="http://127.0.0.1:8000", credential_file=credential(tmp_path)
|
||||
)
|
||||
with pytest.raises(EmbeddingProviderError):
|
||||
await adapter.embed(["one"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_adapter_rejects_profile_that_does_not_match_configured_capability(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
item = profile()
|
||||
async with httpx.AsyncClient(follow_redirects=False) as async_client:
|
||||
with pytest.raises(ValueError, match="capability does not match"):
|
||||
ModelForgeEmbeddingAdapter(
|
||||
item,
|
||||
async_client,
|
||||
base_url="http://modelforge:8000",
|
||||
credential_file=credential(tmp_path),
|
||||
capability="other.embedding@1",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_normalization_uses_the_m5_bfloat16_contract_tolerance(tmp_path: Path) -> None:
|
||||
item = profile()
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
text = json.loads(request.content)["input"][0]
|
||||
vector = [1.019, 0.0, 0.0] if text == "inside" else [1.021, 0.0, 0.0]
|
||||
return httpx.Response(200, json=response(item, [vector]))
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
||||
adapter = ModelForgeEmbeddingAdapter(
|
||||
item,
|
||||
client,
|
||||
base_url="http://127.0.0.1:8000",
|
||||
credential_file=credential(tmp_path),
|
||||
)
|
||||
assert await adapter.embed(["inside"]) == ((1.019, 0.0, 0.0),)
|
||||
with pytest.raises(EmbeddingProviderError, match="not normalized"):
|
||||
await adapter.embed(["outside"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_transient_but_not_unauthorized(tmp_path: Path) -> None:
|
||||
item = profile()
|
||||
attempts = 0
|
||||
|
||||
def transient(_request: httpx.Request) -> httpx.Response:
|
||||
nonlocal attempts
|
||||
attempts += 1
|
||||
return httpx.Response(503 if attempts == 1 else 200, json=response(item, [[1, 0, 0]]))
|
||||
|
||||
async def no_sleep(_seconds: float) -> None:
|
||||
return None
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(transient)) as client:
|
||||
adapter = ModelForgeEmbeddingAdapter(
|
||||
item,
|
||||
client,
|
||||
base_url="http://127.0.0.1:8000",
|
||||
credential_file=credential(tmp_path),
|
||||
sleep=no_sleep,
|
||||
)
|
||||
assert await adapter.embed(["one"]) == ((1.0, 0.0, 0.0),)
|
||||
assert attempts == 2
|
||||
|
||||
denied_attempts = 0
|
||||
|
||||
def denied(_request: httpx.Request) -> httpx.Response:
|
||||
nonlocal denied_attempts
|
||||
denied_attempts += 1
|
||||
return httpx.Response(401)
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(denied)) as client:
|
||||
adapter = ModelForgeEmbeddingAdapter(
|
||||
item,
|
||||
client,
|
||||
base_url="http://127.0.0.1:8000",
|
||||
credential_file=credential(tmp_path),
|
||||
sleep=no_sleep,
|
||||
)
|
||||
with pytest.raises(EmbeddingProviderError, match="credential"):
|
||||
await adapter.embed(["one"])
|
||||
assert denied_attempts == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remote_cloud_host_is_rejected(tmp_path: Path) -> None:
|
||||
async with httpx.AsyncClient() as client:
|
||||
with pytest.raises(ValueError, match="local/private"):
|
||||
ModelForgeEmbeddingAdapter(
|
||||
profile(),
|
||||
client,
|
||||
base_url="https://api.example.com",
|
||||
credential_file=credential(tmp_path),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_successful_batch_gateway_ids_are_retained_in_order(tmp_path: Path) -> None:
|
||||
item = profile()
|
||||
identifiers = [uuid4(), uuid4()]
|
||||
calls = 0
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
nonlocal calls
|
||||
count = len(json.loads(request.content)["input"])
|
||||
body = response(item, [[1, 0, 0]] * count)
|
||||
body["request_id"] = str(identifiers[calls])
|
||||
calls += 1
|
||||
return httpx.Response(200, json=body)
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
||||
adapter = ModelForgeEmbeddingAdapter(
|
||||
item, client, base_url="http://modelforge:8000", credential_file=credential(tmp_path)
|
||||
)
|
||||
vectors, evidence = await adapter.embed_with_evidence(["one", "two", "three"])
|
||||
assert len(vectors) == 3
|
||||
assert evidence.gateway_request_ids == tuple(identifiers)
|
||||
empty, empty_evidence = await adapter.embed_with_evidence([])
|
||||
assert empty == () and empty_evidence.gateway_request_ids == ()
|
||||
assert calls == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"identifier", [None, "not-a-uuid", "00000000-0000-0000-0000-000000000000", 12, True]
|
||||
)
|
||||
async def test_missing_or_invalid_gateway_identity_fails_closed(
|
||||
tmp_path: Path, identifier: object
|
||||
) -> None:
|
||||
item = profile()
|
||||
calls = 0
|
||||
|
||||
def handler(_request: httpx.Request) -> httpx.Response:
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
body = response(item, [[1, 0, 0]])
|
||||
body["request_id"] = identifier
|
||||
return httpx.Response(200, json=body)
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
||||
adapter = ModelForgeEmbeddingAdapter(
|
||||
item, client, base_url="http://modelforge:8000", credential_file=credential(tmp_path)
|
||||
)
|
||||
with pytest.raises(EmbeddingProviderError, match="gateway request identity"):
|
||||
await adapter.embed_with_evidence(["one"])
|
||||
assert calls == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duplicate_gateway_identity_across_batches_is_rejected(tmp_path: Path) -> None:
|
||||
item = profile()
|
||||
identifier = uuid4()
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
count = len(json.loads(request.content)["input"])
|
||||
body = response(item, [[1, 0, 0]] * count)
|
||||
body["request_id"] = str(identifier)
|
||||
return httpx.Response(200, json=body)
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
||||
adapter = ModelForgeEmbeddingAdapter(
|
||||
item, client, base_url="http://modelforge:8000", credential_file=credential(tmp_path)
|
||||
)
|
||||
with pytest.raises(EmbeddingProviderError, match="gateway request identity"):
|
||||
await adapter.embed_with_evidence(["one", "two", "three"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_only_successful_retry_contributes_gateway_evidence(tmp_path: Path) -> None:
|
||||
item = profile()
|
||||
identifiers = [uuid4(), uuid4()]
|
||||
calls = 0
|
||||
|
||||
def handler(_request: httpx.Request) -> httpx.Response:
|
||||
nonlocal calls
|
||||
body = response(item, [[1, 0, 0]])
|
||||
body["request_id"] = str(identifiers[calls])
|
||||
calls += 1
|
||||
return httpx.Response(503 if calls == 1 else 200, json=body)
|
||||
|
||||
async def no_sleep(_seconds: float) -> None:
|
||||
return None
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
||||
adapter = ModelForgeEmbeddingAdapter(
|
||||
item,
|
||||
client,
|
||||
base_url="http://modelforge:8000",
|
||||
credential_file=credential(tmp_path),
|
||||
sleep=no_sleep,
|
||||
)
|
||||
_, evidence = await adapter.embed_with_evidence(["unchanged exact query"], query=True)
|
||||
assert evidence.gateway_request_ids == (identifiers[1],)
|
||||
assert calls == 2
|
||||
@@ -0,0 +1,112 @@
|
||||
"""Real Postgres coverage for the generation profile registry store."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
||||
|
||||
from ragcore.config import Settings
|
||||
from ragcore.domain.control.entities import DataPolicy
|
||||
from ragcore.domain.models import GenerationProfile, GenerationProvider
|
||||
from ragcore.infrastructure.database import create_database_engine, create_session_factory
|
||||
from ragcore.infrastructure.database.models.generation import (
|
||||
GenerationProfileRepositoryError,
|
||||
PostgresGenerationProfileStore,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
|
||||
def _profile(**overrides: object) -> GenerationProfile:
|
||||
values: dict[str, object] = {
|
||||
"id": uuid4(),
|
||||
"profile": "answer-hybrid",
|
||||
"version": uuid4().hex[:12],
|
||||
"provider": GenerationProvider.OLLAMA,
|
||||
"model_identifier": "llama3.1:8b-instruct-q4_0",
|
||||
"model_digest": "a" * 64,
|
||||
"context_limit": 8192,
|
||||
"max_output_tokens": 1024,
|
||||
"temperature": 0.4,
|
||||
"data_policy": DataPolicy.LOCAL_ONLY,
|
||||
}
|
||||
values.update(overrides)
|
||||
return GenerationProfile(**values)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def db_engine() -> AsyncIterator[AsyncEngine]:
|
||||
engine = create_database_engine(Settings(env="test"))
|
||||
try:
|
||||
yield engine
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def store(db_engine: AsyncEngine) -> AsyncIterator[PostgresGenerationProfileStore]:
|
||||
async with db_engine.connect() as connection:
|
||||
await connection.execute(text("TRUNCATE models.generation_profiles"))
|
||||
await connection.commit()
|
||||
yield PostgresGenerationProfileStore(create_session_factory(db_engine))
|
||||
|
||||
|
||||
async def test_add_and_get_round_trip(store: PostgresGenerationProfileStore) -> None:
|
||||
profile = _profile()
|
||||
await store.add(profile)
|
||||
fetched = await store.get(profile.id)
|
||||
assert fetched == profile
|
||||
|
||||
|
||||
async def test_get_returns_none_for_an_unknown_id(store: PostgresGenerationProfileStore) -> None:
|
||||
assert await store.get(uuid4()) is None
|
||||
|
||||
|
||||
async def test_adding_the_same_id_twice_is_rejected(store: PostgresGenerationProfileStore) -> None:
|
||||
profile = _profile()
|
||||
await store.add(profile)
|
||||
with pytest.raises(GenerationProfileRepositoryError):
|
||||
await store.add(profile)
|
||||
|
||||
|
||||
async def test_adding_the_same_profile_version_twice_is_rejected(
|
||||
store: PostgresGenerationProfileStore,
|
||||
) -> None:
|
||||
profile = _profile()
|
||||
await store.add(profile)
|
||||
with pytest.raises(GenerationProfileRepositoryError):
|
||||
await store.add(_profile(id=uuid4(), profile=profile.profile, version=profile.version))
|
||||
|
||||
|
||||
async def test_activate_requires_a_registered_profile(
|
||||
store: PostgresGenerationProfileStore,
|
||||
) -> None:
|
||||
with pytest.raises(GenerationProfileRepositoryError):
|
||||
await store.activate(uuid4())
|
||||
|
||||
|
||||
async def test_activating_a_second_profile_deactivates_the_first(
|
||||
store: PostgresGenerationProfileStore,
|
||||
) -> None:
|
||||
first = _profile()
|
||||
second = _profile(id=uuid4(), version=uuid4().hex[:12])
|
||||
await store.add(first)
|
||||
await store.add(second)
|
||||
|
||||
await store.activate(first.id)
|
||||
assert (await store.get_active()).id == first.id # type: ignore[union-attr]
|
||||
|
||||
await store.activate(second.id)
|
||||
active = await store.get_active()
|
||||
assert active is not None
|
||||
assert active.id == second.id
|
||||
|
||||
|
||||
async def test_get_active_returns_none_when_nothing_is_activated(
|
||||
store: PostgresGenerationProfileStore,
|
||||
) -> None:
|
||||
assert await store.get_active() is None
|
||||
@@ -0,0 +1,256 @@
|
||||
"""OllamaGenerationAdapter: local-host validation, streaming, retries, breaker."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from uuid import uuid4
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from ragcore.application.limits import ConcurrencyBudget
|
||||
from ragcore.domain.control.entities import DataPolicy
|
||||
from ragcore.domain.models import GenerationProfile, GenerationProvider
|
||||
from ragcore.infrastructure.models.ollama_generation import (
|
||||
CIRCUIT_FAILURE_THRESHOLD,
|
||||
CircuitOpenError,
|
||||
GenerationProviderError,
|
||||
OllamaGenerationAdapter,
|
||||
assemble,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
|
||||
def profile(**overrides: object) -> GenerationProfile:
|
||||
values: dict[str, object] = {
|
||||
"id": uuid4(),
|
||||
"profile": "answer-hybrid",
|
||||
"version": "v1",
|
||||
"provider": GenerationProvider.OLLAMA,
|
||||
"model_identifier": "llama3.1:8b-instruct-q4_0",
|
||||
"model_digest": "a" * 64,
|
||||
"context_limit": 4096,
|
||||
"max_output_tokens": 512,
|
||||
"temperature": 0.2,
|
||||
"data_policy": DataPolicy.LOCAL_ONLY,
|
||||
"timeout_seconds": 5.0,
|
||||
"max_retries": 2,
|
||||
}
|
||||
values.update(overrides)
|
||||
return GenerationProfile(**values)
|
||||
|
||||
|
||||
async def no_sleep(_: float) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def _ndjson(*lines: dict[str, object]) -> bytes:
|
||||
return b"\n".join(json.dumps(line).encode() for line in lines)
|
||||
|
||||
|
||||
def _client(handler) -> httpx.AsyncClient:
|
||||
return httpx.AsyncClient(transport=httpx.MockTransport(handler), follow_redirects=False)
|
||||
|
||||
|
||||
def _adapter(handler, **overrides: object) -> OllamaGenerationAdapter:
|
||||
clock = overrides.pop("clock", None)
|
||||
kwargs: dict[str, object] = {"base_url": "http://ollama", "sleep": no_sleep}
|
||||
if clock is not None:
|
||||
kwargs["clock"] = clock
|
||||
return OllamaGenerationAdapter(profile(**overrides), _client(handler), **kwargs)
|
||||
|
||||
|
||||
def test_adapter_rejects_nonlocal_provider_hosts() -> None:
|
||||
with pytest.raises(ValueError, match="allowlisted local"):
|
||||
OllamaGenerationAdapter(
|
||||
profile(), _client(lambda request: httpx.Response(200)), base_url="http://169.254.169.254"
|
||||
)
|
||||
|
||||
|
||||
async def test_generate_stream_yields_each_chunk_in_order() -> None:
|
||||
body = _ndjson(
|
||||
{"response": "Hel", "done": False},
|
||||
{"response": "lo", "done": False},
|
||||
{"response": "", "done": True},
|
||||
)
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert json.loads(request.content)["stream"] is True
|
||||
return httpx.Response(200, content=body)
|
||||
|
||||
adapter = _adapter(handler)
|
||||
chunks = [chunk async for chunk in adapter.generate_stream("hi")]
|
||||
assert [chunk.text for chunk in chunks] == ["Hel", "lo", ""]
|
||||
assert chunks[-1].done is True
|
||||
|
||||
|
||||
async def test_generate_assembles_the_same_content_a_stream_would() -> None:
|
||||
body = _ndjson(
|
||||
{"response": "The answer is ", "done": False},
|
||||
{"response": "42.", "done": True},
|
||||
)
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, content=body)
|
||||
|
||||
adapter = _adapter(handler)
|
||||
result = await adapter.generate("what is the answer?")
|
||||
streamed = [chunk async for chunk in adapter.generate_stream("what is the answer?")]
|
||||
|
||||
assert result.text == "The answer is 42."
|
||||
assert result == assemble(streamed)
|
||||
|
||||
|
||||
async def test_breaking_out_of_the_stream_early_does_not_raise() -> None:
|
||||
body = _ndjson(
|
||||
{"response": "one", "done": False},
|
||||
{"response": "two", "done": False},
|
||||
{"response": "three", "done": True},
|
||||
)
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, content=body)
|
||||
|
||||
adapter = _adapter(handler)
|
||||
seen = []
|
||||
async for chunk in adapter.generate_stream("hi"):
|
||||
seen.append(chunk.text)
|
||||
if len(seen) == 1:
|
||||
break
|
||||
assert seen == ["one"]
|
||||
|
||||
|
||||
async def test_retryable_status_retries_before_any_streaming_then_succeeds() -> None:
|
||||
attempts = {"count": 0}
|
||||
body = _ndjson({"response": "ok", "done": True})
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
attempts["count"] += 1
|
||||
if attempts["count"] < 3:
|
||||
return httpx.Response(503)
|
||||
return httpx.Response(200, content=body)
|
||||
|
||||
adapter = _adapter(handler)
|
||||
result = await adapter.generate("hi")
|
||||
assert result.text == "ok"
|
||||
assert attempts["count"] == 3
|
||||
|
||||
|
||||
async def test_exhausting_retries_raises_generation_unavailable() -> None:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(503)
|
||||
|
||||
adapter = _adapter(handler, max_retries=1)
|
||||
with pytest.raises(GenerationProviderError, match="generation_unavailable"):
|
||||
await adapter.generate("hi")
|
||||
|
||||
|
||||
async def test_non_retryable_status_raises_generation_rejected() -> None:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(400)
|
||||
|
||||
adapter = _adapter(handler)
|
||||
with pytest.raises(GenerationProviderError, match="generation_rejected"):
|
||||
await adapter.generate("hi")
|
||||
|
||||
|
||||
async def test_invalid_response_line_is_a_stable_error() -> None:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, content=b"not json at all")
|
||||
|
||||
adapter = _adapter(handler)
|
||||
with pytest.raises(GenerationProviderError, match="generation_response_invalid"):
|
||||
await adapter.generate("hi")
|
||||
|
||||
|
||||
async def test_stream_ending_without_a_done_line_is_a_stable_error() -> None:
|
||||
body = _ndjson({"response": "partial", "done": False})
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, content=body)
|
||||
|
||||
adapter = _adapter(handler)
|
||||
with pytest.raises(GenerationProviderError, match="generation_response_invalid"):
|
||||
await adapter.generate("hi")
|
||||
|
||||
|
||||
async def test_empty_and_oversized_prompts_are_rejected() -> None:
|
||||
adapter = _adapter(lambda request: httpx.Response(200))
|
||||
with pytest.raises(ValueError, match="empty"):
|
||||
await adapter.generate("")
|
||||
with pytest.raises(ValueError, match="bounded"):
|
||||
await adapter.generate("x" * 200_001)
|
||||
|
||||
|
||||
async def test_circuit_opens_after_consecutive_failures_and_blocks_further_calls() -> None:
|
||||
attempts = {"count": 0}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
attempts["count"] += 1
|
||||
return httpx.Response(503)
|
||||
|
||||
adapter = _adapter(handler, max_retries=0)
|
||||
for _ in range(CIRCUIT_FAILURE_THRESHOLD):
|
||||
with pytest.raises(GenerationProviderError, match="generation_unavailable"):
|
||||
await adapter.generate("hi")
|
||||
|
||||
calls_before_circuit_check = attempts["count"]
|
||||
with pytest.raises(CircuitOpenError):
|
||||
await adapter.generate("hi")
|
||||
assert attempts["count"] == calls_before_circuit_check # the breaker skipped the network call
|
||||
|
||||
|
||||
async def test_circuit_half_opens_after_cooldown_and_closes_on_success() -> None:
|
||||
outcomes = iter(
|
||||
[httpx.Response(503)] * CIRCUIT_FAILURE_THRESHOLD
|
||||
+ [httpx.Response(200, content=_ndjson({"response": "ok", "done": True}))]
|
||||
)
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return next(outcomes)
|
||||
|
||||
clock_time = {"value": 0.0}
|
||||
|
||||
def clock() -> float:
|
||||
return clock_time["value"]
|
||||
|
||||
adapter = _adapter(handler, max_retries=0, clock=clock)
|
||||
for _ in range(CIRCUIT_FAILURE_THRESHOLD):
|
||||
with pytest.raises(GenerationProviderError, match="generation_unavailable"):
|
||||
await adapter.generate("hi")
|
||||
|
||||
# Still within the cooldown window: the breaker rejects without a network call.
|
||||
with pytest.raises(CircuitOpenError):
|
||||
await adapter.generate("hi")
|
||||
|
||||
# Advance past the cooldown: the next call is a real, half-open attempt.
|
||||
clock_time["value"] += 3600.0
|
||||
result = await adapter.generate("hi")
|
||||
assert result.text == "ok"
|
||||
|
||||
|
||||
async def test_adapter_surfaces_concurrency_exhaustion_as_a_provider_error() -> None:
|
||||
budget = ConcurrencyBudget(name="generation-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 = OllamaGenerationAdapter(
|
||||
profile(),
|
||||
_client(lambda request: httpx.Response(200)),
|
||||
base_url="http://ollama",
|
||||
sleep=no_sleep,
|
||||
concurrency_budget=budget,
|
||||
)
|
||||
with pytest.raises(GenerationProviderError, match="generation_concurrency_exhausted"):
|
||||
await adapter.generate("hi")
|
||||
|
||||
release.set()
|
||||
await holder
|
||||
@@ -0,0 +1,170 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from ragcore.infrastructure.models import (
|
||||
ModelForgeRerankingAdapter,
|
||||
RerankDocument,
|
||||
)
|
||||
from ragcore.retrieval.rerank.models import RerankProviderError
|
||||
|
||||
|
||||
def credential(tmp_path: Path) -> Path:
|
||||
path = tmp_path / "reranking-credential"
|
||||
path.write_text("mf_" + "x" * 48, encoding="utf-8")
|
||||
path.chmod(0o600)
|
||||
return path
|
||||
|
||||
|
||||
def response(ids: list[str]) -> dict[str, object]:
|
||||
return {
|
||||
"capability": "rag.reranking@1",
|
||||
"results": [
|
||||
{"id": identifier, "score": 1.0 / rank, "rank": rank}
|
||||
for rank, identifier in enumerate(ids, 1)
|
||||
],
|
||||
"request_id": str(uuid4()),
|
||||
"execution": {
|
||||
"cold": False,
|
||||
"node": "Tower",
|
||||
"residency": "warm",
|
||||
"load_count": 1,
|
||||
"timings": {
|
||||
"gateway_ms": 1,
|
||||
"queue_ms": 0,
|
||||
"load_ms": 0,
|
||||
"inference_ms": 2,
|
||||
"total_ms": 3,
|
||||
},
|
||||
},
|
||||
"usage": {"input_count": len(ids), "input_tokens": 10},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scoped_capability_request_and_ranked_id_validation(tmp_path: Path) -> None:
|
||||
seen: list[httpx.Request] = []
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
seen.append(request)
|
||||
body = json.loads(request.content)
|
||||
return httpx.Response(200, json=response([body["documents"][1]["id"]]))
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
||||
adapter = ModelForgeRerankingAdapter(
|
||||
client,
|
||||
base_url="http://modelforge:8000",
|
||||
credential_file=credential(tmp_path),
|
||||
)
|
||||
ranked, evidence = await adapter.rerank(
|
||||
"which passage?",
|
||||
(RerankDocument("one", "irrelevant"), RerankDocument("two", "relevant")),
|
||||
top_n=1,
|
||||
)
|
||||
assert [(item.id, item.rank) for item in ranked] == [("two", 1)]
|
||||
assert seen[0].url.path == "/api/v1/capabilities/rag.reranking@1/invoke"
|
||||
assert seen[0].headers["authorization"].startswith("Bearer mf_")
|
||||
assert seen[0].headers["x-correlation-id"].startswith("ragcore-rerank-")
|
||||
assert evidence.gateway_round_trip_ms >= 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("status", [401, 403])
|
||||
async def test_scope_or_auth_rejection_fails_closed(tmp_path: Path, status: int) -> None:
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.MockTransport(lambda _request: httpx.Response(status))
|
||||
) as client:
|
||||
adapter = ModelForgeRerankingAdapter(
|
||||
client,
|
||||
base_url="http://127.0.0.1:8000",
|
||||
credential_file=credential(tmp_path),
|
||||
)
|
||||
with pytest.raises(RerankProviderError, match="scope_rejected"):
|
||||
await adapter.rerank("query", (RerankDocument("one", "text"),), top_n=1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("failure", ["duplicate", "unknown", "rank", "nan"])
|
||||
async def test_invalid_ranking_response_fails_closed(tmp_path: Path, failure: str) -> None:
|
||||
body = response(["one", "two"])
|
||||
results = body["results"]
|
||||
assert isinstance(results, list)
|
||||
if failure == "duplicate":
|
||||
results[1]["id"] = "one"
|
||||
elif failure == "unknown":
|
||||
results[1]["id"] = "three"
|
||||
elif failure == "rank":
|
||||
results[1]["rank"] = 1
|
||||
else:
|
||||
results[1]["score"] = float("nan")
|
||||
|
||||
def handler(_request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
200,
|
||||
content=json.dumps(body, allow_nan=True).encode(),
|
||||
headers={"content-type": "application/json"},
|
||||
)
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
||||
adapter = ModelForgeRerankingAdapter(
|
||||
client,
|
||||
base_url="http://127.0.0.1:8000",
|
||||
credential_file=credential(tmp_path),
|
||||
)
|
||||
with pytest.raises(RerankProviderError, match="result_invalid"):
|
||||
await adapter.rerank(
|
||||
"query",
|
||||
(RerankDocument("one", "a"), RerankDocument("two", "b")),
|
||||
top_n=2,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_input_bounds_apply_before_network(tmp_path: Path) -> None:
|
||||
calls = 0
|
||||
|
||||
def handler(_request: httpx.Request) -> httpx.Response:
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
return httpx.Response(500)
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
||||
adapter = ModelForgeRerankingAdapter(
|
||||
client,
|
||||
base_url="http://127.0.0.1:8000",
|
||||
credential_file=credential(tmp_path),
|
||||
)
|
||||
with pytest.raises(ValueError, match="document count"):
|
||||
await adapter.rerank(
|
||||
"query", tuple(RerankDocument(str(index), "text") for index in range(41))
|
||||
)
|
||||
with pytest.raises(ValueError, match="unique"):
|
||||
await adapter.rerank(
|
||||
"query",
|
||||
(RerankDocument("same", "a"), RerankDocument("same", "b")),
|
||||
top_n=2,
|
||||
)
|
||||
assert calls == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remote_host_and_redirecting_client_are_rejected(tmp_path: Path) -> None:
|
||||
async with httpx.AsyncClient() as client:
|
||||
with pytest.raises(ValueError, match="local/private"):
|
||||
ModelForgeRerankingAdapter(
|
||||
client,
|
||||
base_url="https://api.example.com",
|
||||
credential_file=credential(tmp_path),
|
||||
)
|
||||
async with httpx.AsyncClient(follow_redirects=True) as client:
|
||||
with pytest.raises(ValueError, match="redirects"):
|
||||
ModelForgeRerankingAdapter(
|
||||
client,
|
||||
base_url="http://127.0.0.1:8000",
|
||||
credential_file=credential(tmp_path),
|
||||
)
|
||||
Reference in new issue
Block a user