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,219 @@
|
||||
"""OllamaRerankAdapter: local-host validation, retries, and response-shape checks."""
|
||||
|
||||
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 import DataPolicy
|
||||
from ragcore.infrastructure.models.reranker import RerankerProfile, RerankerProvider
|
||||
from ragcore.infrastructure.models.reranker.ollama import OllamaRerankAdapter
|
||||
from ragcore.retrieval.rerank import RerankProviderError
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
|
||||
def profile(**changes: object) -> RerankerProfile:
|
||||
value: dict[str, object] = {
|
||||
"id": uuid4(),
|
||||
"profile": "local",
|
||||
"version": "v1",
|
||||
"provider": RerankerProvider.OLLAMA,
|
||||
"model_identifier": "bge-reranker-v2-m3",
|
||||
"model_digest": "a" * 64,
|
||||
"context_limit": 512,
|
||||
"data_policy": DataPolicy.LOCAL_ONLY,
|
||||
"timeout_seconds": 1.0,
|
||||
"max_candidates": 4,
|
||||
"max_retries": 2,
|
||||
}
|
||||
value.update(changes)
|
||||
return RerankerProfile(**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"):
|
||||
OllamaRerankAdapter(
|
||||
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 = "bge-reranker:latest" # type: ignore[misc]
|
||||
with pytest.raises(ValueError, match="floating"):
|
||||
profile(model_identifier="bge-reranker:latest")
|
||||
with pytest.raises(ValueError, match="digest"):
|
||||
profile(model_digest="not-a-digest")
|
||||
with pytest.raises(ValueError, match="context_limit"):
|
||||
profile(context_limit=0)
|
||||
|
||||
|
||||
async def test_adapter_sends_query_and_documents_and_reorders_by_index() -> None:
|
||||
client = FakeClient(
|
||||
[
|
||||
response(
|
||||
200,
|
||||
{
|
||||
"results": [
|
||||
{"index": 1, "relevance_score": 0.9},
|
||||
{"index": 0, "relevance_score": 0.1},
|
||||
]
|
||||
},
|
||||
)
|
||||
]
|
||||
)
|
||||
adapter = OllamaRerankAdapter(profile(), client, base_url="http://ollama", sleep=no_sleep) # type: ignore[arg-type]
|
||||
|
||||
scores = await adapter.score("a query", ("doc one", "doc two"))
|
||||
|
||||
assert scores == (0.1, 0.9)
|
||||
assert client.requests[0]["json"] == {
|
||||
"model": "bge-reranker-v2-m3",
|
||||
"query": "a query",
|
||||
"documents": ["doc one", "doc two"],
|
||||
}
|
||||
|
||||
|
||||
async def test_adapter_retries_on_retryable_status_then_succeeds() -> None:
|
||||
client = FakeClient(
|
||||
[
|
||||
response(503, {"error": "busy"}),
|
||||
response(429, {"error": "busy"}),
|
||||
response(200, {"results": [{"index": 0, "relevance_score": 0.5}]}),
|
||||
]
|
||||
)
|
||||
adapter = OllamaRerankAdapter(profile(), client, base_url="http://ollama", sleep=no_sleep) # type: ignore[arg-type]
|
||||
|
||||
assert await adapter.score("q", ("one",)) == (0.5,)
|
||||
assert len(client.requests) == 3
|
||||
|
||||
|
||||
async def test_adapter_raises_after_exhausting_retries() -> None:
|
||||
client = FakeClient([response(503, {}), response(503, {}), response(503, {})])
|
||||
adapter = OllamaRerankAdapter(
|
||||
profile(max_retries=2), client, base_url="http://ollama", sleep=no_sleep # type: ignore[arg-type]
|
||||
)
|
||||
with pytest.raises(RerankProviderError, match="rerank_unavailable"):
|
||||
await adapter.score("q", ("one",))
|
||||
|
||||
|
||||
async def test_adapter_rejects_cardinality_mismatch_and_duplicate_index() -> None:
|
||||
mismatch = OllamaRerankAdapter(
|
||||
profile(),
|
||||
FakeClient([response(200, {"results": [{"index": 0, "relevance_score": 0.5}]})]), # type: ignore[arg-type]
|
||||
base_url="http://ollama",
|
||||
sleep=no_sleep,
|
||||
)
|
||||
with pytest.raises(RerankProviderError, match="rerank_cardinality_mismatch"):
|
||||
await mismatch.score("q", ("one", "two"))
|
||||
|
||||
duplicate = OllamaRerankAdapter(
|
||||
profile(),
|
||||
FakeClient(
|
||||
[
|
||||
response(
|
||||
200,
|
||||
{
|
||||
"results": [
|
||||
{"index": 0, "relevance_score": 0.1},
|
||||
{"index": 0, "relevance_score": 0.9},
|
||||
]
|
||||
},
|
||||
)
|
||||
]
|
||||
), # type: ignore[arg-type]
|
||||
base_url="http://ollama",
|
||||
sleep=no_sleep,
|
||||
)
|
||||
with pytest.raises(RerankProviderError, match="rerank_result_duplicate_index"):
|
||||
await duplicate.score("q", ("one", "two"))
|
||||
|
||||
|
||||
async def test_adapter_rejects_nonfinite_score() -> None:
|
||||
# json.dumps allows NaN by default (unlike httpx's stricter json= encoder), so a raw
|
||||
# content body is used here to construct an out-of-spec response body deterministically.
|
||||
body = json.dumps({"results": [{"index": 0, "relevance_score": float("nan")}]}).encode()
|
||||
non_finite_response = httpx.Response(
|
||||
200, content=body, request=httpx.Request("POST", "http://test")
|
||||
)
|
||||
adapter = OllamaRerankAdapter(
|
||||
profile(),
|
||||
FakeClient([non_finite_response]), # type: ignore[arg-type]
|
||||
base_url="http://ollama",
|
||||
sleep=no_sleep,
|
||||
)
|
||||
with pytest.raises(RerankProviderError, match="rerank_result_score_invalid"):
|
||||
await adapter.score("q", ("one",))
|
||||
|
||||
|
||||
async def test_adapter_rejects_input_exceeding_profile_bounds() -> None:
|
||||
adapter = OllamaRerankAdapter(
|
||||
profile(max_candidates=1), FakeClient([]), base_url="http://ollama", sleep=no_sleep # type: ignore[arg-type]
|
||||
)
|
||||
with pytest.raises(ValueError, match="candidate limit"):
|
||||
await adapter.score("q", ("one", "two"))
|
||||
|
||||
empty_query = OllamaRerankAdapter(
|
||||
profile(), FakeClient([]), base_url="http://ollama", sleep=no_sleep # type: ignore[arg-type]
|
||||
)
|
||||
with pytest.raises(ValueError, match="query"):
|
||||
await empty_query.score("", ("one",))
|
||||
|
||||
|
||||
async def test_adapter_returns_empty_tuple_for_no_documents() -> None:
|
||||
adapter = OllamaRerankAdapter(
|
||||
profile(), FakeClient([]), base_url="http://ollama", sleep=no_sleep # type: ignore[arg-type]
|
||||
)
|
||||
assert await adapter.score("q", ()) == ()
|
||||
|
||||
|
||||
async def test_adapter_surfaces_concurrency_exhaustion_as_a_provider_error() -> None:
|
||||
budget = ConcurrencyBudget(name="rerank-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 = OllamaRerankAdapter(
|
||||
profile(),
|
||||
FakeClient([]), # type: ignore[arg-type]
|
||||
base_url="http://ollama",
|
||||
sleep=no_sleep,
|
||||
concurrency_budget=budget,
|
||||
)
|
||||
with pytest.raises(RerankProviderError, match="rerank_concurrency_exhausted"):
|
||||
await adapter.score("q", ("one",))
|
||||
|
||||
release.set()
|
||||
await holder
|
||||
@@ -0,0 +1,188 @@
|
||||
"""RerankService: bounded reranking, deterministic ordering and degraded fallback."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Sequence
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
from ragcore.domain.control import DataPolicy
|
||||
from ragcore.infrastructure.models.reranker import RerankerProfile, RerankerProvider
|
||||
from ragcore.retrieval.ranking.rrf import FusedCandidate
|
||||
from ragcore.retrieval.rerank import (
|
||||
MAX_RERANK_CANDIDATES,
|
||||
RerankCandidateInput,
|
||||
RerankOutcomeState,
|
||||
RerankProviderError,
|
||||
RerankService,
|
||||
)
|
||||
from ragcore.retrieval.search.models import Candidate, CandidateBackendKind, CandidateProvenance
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
|
||||
def reranker_profile(**changes: object) -> RerankerProfile:
|
||||
value: dict[str, object] = {
|
||||
"id": uuid4(),
|
||||
"profile": "local",
|
||||
"version": "v1",
|
||||
"provider": RerankerProvider.OLLAMA,
|
||||
"model_identifier": "bge-reranker-v2-m3",
|
||||
"model_digest": "a" * 64,
|
||||
"context_limit": 512,
|
||||
"data_policy": DataPolicy.LOCAL_ONLY,
|
||||
"timeout_seconds": 1.0,
|
||||
"max_candidates": 40,
|
||||
"max_retries": 1,
|
||||
}
|
||||
value.update(changes)
|
||||
return RerankerProfile(**value)
|
||||
|
||||
|
||||
def _candidate(index: int) -> Candidate:
|
||||
chunk_id = uuid4()
|
||||
return Candidate(
|
||||
point_id=uuid4(),
|
||||
score=1.0 / (index + 1),
|
||||
rank=index + 1,
|
||||
backend=CandidateBackendKind.DENSE,
|
||||
provenance=CandidateProvenance(
|
||||
workspace_id=uuid4(),
|
||||
space_id=uuid4(),
|
||||
document_id=uuid4(),
|
||||
document_version_id=uuid4(),
|
||||
chunk_id=chunk_id,
|
||||
source_id=uuid4(),
|
||||
parent_chunk_id=None,
|
||||
status="published",
|
||||
content_type="text/markdown",
|
||||
language="en",
|
||||
trust_level="source",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _fused(index: int) -> FusedCandidate:
|
||||
return FusedCandidate(
|
||||
candidate=_candidate(index),
|
||||
rrf_score=1.0 / (60 + index + 1),
|
||||
dense_rank=index + 1,
|
||||
sparse_rank=None,
|
||||
)
|
||||
|
||||
|
||||
def _inputs(count: int) -> tuple[RerankCandidateInput, ...]:
|
||||
return tuple(
|
||||
RerankCandidateInput(fused=_fused(index), text=f"passage {index}") for index in range(count)
|
||||
)
|
||||
|
||||
|
||||
class FixtureReranker:
|
||||
"""Deterministic fixture: score is simply the reverse of input order."""
|
||||
|
||||
def __init__(self, scores: Sequence[float] | None = None) -> None:
|
||||
self._scores = scores
|
||||
|
||||
async def score(self, query: str, documents: Sequence[str]) -> tuple[float, ...]:
|
||||
if self._scores is not None:
|
||||
return tuple(self._scores)
|
||||
count = len(documents)
|
||||
return tuple(float(count - index) for index in range(count))
|
||||
|
||||
|
||||
class HangingReranker:
|
||||
async def score(self, query: str, documents: Sequence[str]) -> tuple[float, ...]:
|
||||
await asyncio.sleep(10)
|
||||
return tuple(0.0 for _ in documents)
|
||||
|
||||
|
||||
class FailingReranker:
|
||||
def __init__(self, code: str) -> None:
|
||||
self._code = code
|
||||
|
||||
async def score(self, query: str, documents: Sequence[str]) -> tuple[float, ...]:
|
||||
raise RerankProviderError(self._code)
|
||||
|
||||
|
||||
async def test_rerank_reorders_by_score_deterministically() -> None:
|
||||
service = RerankService(reranker_profile(), FixtureReranker())
|
||||
result = await service.rerank("query", _inputs(5))
|
||||
|
||||
assert result.state is RerankOutcomeState.RERANKED
|
||||
assert not result.degraded
|
||||
# FixtureReranker scores highest for the first document, so order is preserved.
|
||||
assert [item.fused.candidate.rank for item in result.candidates] == [1, 2, 3, 4, 5]
|
||||
assert [item.rerank_rank for item in result.candidates] == [1, 2, 3, 4, 5]
|
||||
assert all(item.rerank_score is not None for item in result.candidates)
|
||||
|
||||
|
||||
async def test_rerank_ties_break_deterministically_by_input_order() -> None:
|
||||
service = RerankService(reranker_profile(), FixtureReranker(scores=[1.0, 1.0, 1.0]))
|
||||
result = await service.rerank("query", _inputs(3))
|
||||
|
||||
assert [item.fused.candidate.rank for item in result.candidates] == [1, 2, 3]
|
||||
|
||||
|
||||
async def test_rerank_enforces_result_limit() -> None:
|
||||
service = RerankService(reranker_profile(), FixtureReranker())
|
||||
result = await service.rerank("query", _inputs(8), result_limit=3)
|
||||
|
||||
assert len(result.candidates) == 3
|
||||
assert [item.fused.candidate.rank for item in result.candidates] == [1, 2, 3]
|
||||
|
||||
|
||||
async def test_rerank_rejects_candidate_input_over_the_bound() -> None:
|
||||
service = RerankService(reranker_profile(), FixtureReranker())
|
||||
with pytest.raises(ValueError, match="bounded limit"):
|
||||
await service.rerank("query", _inputs(MAX_RERANK_CANDIDATES + 1))
|
||||
|
||||
|
||||
async def test_rerank_rejects_result_limit_outside_bound() -> None:
|
||||
service = RerankService(reranker_profile(), FixtureReranker())
|
||||
with pytest.raises(ValueError, match="result_limit"):
|
||||
await service.rerank("query", _inputs(3), result_limit=0)
|
||||
|
||||
|
||||
async def test_rerank_timeout_leaves_fused_order_with_degraded_flag() -> None:
|
||||
service = RerankService(reranker_profile(timeout_seconds=0.1), HangingReranker())
|
||||
result = await service.rerank("query", _inputs(5), result_limit=3)
|
||||
|
||||
assert result.state is RerankOutcomeState.DEGRADED
|
||||
assert result.degraded
|
||||
assert result.degradation_code == "rerank_timeout"
|
||||
assert len(result.candidates) == 3
|
||||
assert [item.fused.candidate.rank for item in result.candidates] == [1, 2, 3]
|
||||
assert all(item.rerank_score is None for item in result.candidates)
|
||||
|
||||
|
||||
async def test_rerank_provider_failure_degrades_with_its_error_code() -> None:
|
||||
service = RerankService(reranker_profile(), FailingReranker("rerank_unavailable"))
|
||||
result = await service.rerank("query", _inputs(4))
|
||||
|
||||
assert result.degraded
|
||||
assert result.degradation_code == "rerank_unavailable"
|
||||
assert [item.fused.candidate.rank for item in result.candidates] == [1, 2, 3, 4]
|
||||
|
||||
|
||||
async def test_rerank_malformed_response_degrades() -> None:
|
||||
wrong_count = RerankService(reranker_profile(), FixtureReranker(scores=[1.0, 2.0]))
|
||||
result = await wrong_count.rerank("query", _inputs(3))
|
||||
assert result.degraded
|
||||
assert result.degradation_code == "rerank_response_invalid"
|
||||
|
||||
non_finite = RerankService(
|
||||
reranker_profile(), FixtureReranker(scores=[float("nan"), 1.0, 2.0])
|
||||
)
|
||||
result = await non_finite.rerank("query", _inputs(3))
|
||||
assert result.degraded
|
||||
assert result.degradation_code == "rerank_response_invalid"
|
||||
|
||||
|
||||
async def test_rerank_empty_candidates_returns_empty_reranked_result() -> None:
|
||||
service = RerankService(reranker_profile(), FixtureReranker())
|
||||
result = await service.rerank("query", ())
|
||||
|
||||
assert result.candidates == ()
|
||||
assert result.state is RerankOutcomeState.RERANKED
|
||||
Reference in new issue
Block a user