Initial curated RAGcore source from c16667f0c9397e2bfc8dd0255b114fba23c8d05e
This commit is contained in:
commit
7ca726caf2
659 files changed
+98171
No files matched your search
@@ -0,0 +1,445 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from qdrant_client.http import models
|
||||
|
||||
from ragcore.application.retrieval import CandidateSearchRequest, CandidateSearchService
|
||||
from ragcore.application.retrieval.authorization import AuthorizedRetrievalScope
|
||||
from ragcore.infrastructure.qdrant.filters import compile_retrieval_filter
|
||||
from ragcore.infrastructure.qdrant.models import CollectionKind, CollectionProfile
|
||||
from ragcore.retrieval.search import (
|
||||
BackendSearchRequest,
|
||||
Candidate,
|
||||
CandidateBackendError,
|
||||
CandidateBackendKind,
|
||||
CandidateBackendState,
|
||||
CandidateProvenance,
|
||||
DenseQuery,
|
||||
EvaluationRetrievalTarget,
|
||||
QdrantCandidateBackend,
|
||||
)
|
||||
from ragcore.retrieval.sparse import SparseEncoder
|
||||
|
||||
WORKSPACE_ID = uuid4()
|
||||
SPACE_ID = uuid4()
|
||||
|
||||
|
||||
PROFILE = CollectionProfile(
|
||||
kind=CollectionKind.DENSE,
|
||||
profile="hybrid",
|
||||
version="v1",
|
||||
dimensions=2,
|
||||
model_digest="0" * 64,
|
||||
)
|
||||
|
||||
|
||||
def candidate(
|
||||
backend: CandidateBackendKind,
|
||||
rank: int = 1,
|
||||
*,
|
||||
workspace_id: object = WORKSPACE_ID,
|
||||
space_id: object = SPACE_ID,
|
||||
) -> Candidate:
|
||||
return Candidate(
|
||||
point_id=uuid4(),
|
||||
score=1.0 / rank,
|
||||
rank=rank,
|
||||
backend=backend,
|
||||
provenance=CandidateProvenance(
|
||||
workspace_id=workspace_id,
|
||||
space_id=space_id,
|
||||
document_id=uuid4(),
|
||||
document_version_id=uuid4(),
|
||||
chunk_id=uuid4(),
|
||||
source_id=uuid4(),
|
||||
parent_chunk_id=None,
|
||||
status="published",
|
||||
content_type="text/plain",
|
||||
language="en",
|
||||
trust_level="source",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class RecordingBackend:
|
||||
def __init__(self, result: tuple[Candidate, ...] = ()) -> None:
|
||||
self.result = result
|
||||
self.requests: list[BackendSearchRequest] = []
|
||||
|
||||
async def search(self, request: BackendSearchRequest) -> tuple[Candidate, ...]:
|
||||
self.requests.append(request)
|
||||
return self.result
|
||||
|
||||
|
||||
class FailingBackend:
|
||||
async def search(self, request: BackendSearchRequest) -> tuple[Candidate, ...]:
|
||||
raise CandidateBackendError("backend_unavailable")
|
||||
|
||||
|
||||
class SlowBackend:
|
||||
async def search(self, request: BackendSearchRequest) -> tuple[Candidate, ...]:
|
||||
await asyncio.sleep(1)
|
||||
return ()
|
||||
|
||||
|
||||
def request(*, timeout_seconds: float = 1.0, dense_limit: int = 50) -> CandidateSearchRequest:
|
||||
return CandidateSearchRequest(
|
||||
profile=PROFILE,
|
||||
dense_query=DenseQuery((0.1, 0.2)),
|
||||
sparse_query=SparseEncoder().encode_query("HTTP_500"),
|
||||
filter_plan=compile_retrieval_filter(
|
||||
AuthorizedRetrievalScope(WORKSPACE_ID, frozenset({SPACE_ID}))
|
||||
),
|
||||
dense_limit=dense_limit,
|
||||
timeout_seconds=timeout_seconds,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dense_and_sparse_run_in_parallel_and_preserve_provenance() -> None:
|
||||
dense = RecordingBackend((candidate(CandidateBackendKind.DENSE),))
|
||||
sparse = RecordingBackend((candidate(CandidateBackendKind.SPARSE),))
|
||||
|
||||
result = await CandidateSearchService(dense, sparse).search(request())
|
||||
|
||||
assert result.degraded is False
|
||||
assert result.dense.candidates[0].provenance.document_id.int != 0
|
||||
assert result.sparse.candidates[0].provenance.source_id.int != 0
|
||||
assert dense.requests[0].limit == 50
|
||||
assert sparse.requests[0].backend is CandidateBackendKind.SPARSE
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_partial_backend_failure_is_explicit_degraded_metadata() -> None:
|
||||
dense = RecordingBackend((candidate(CandidateBackendKind.DENSE),))
|
||||
|
||||
result = await CandidateSearchService(dense, FailingBackend()).search(request())
|
||||
|
||||
assert result.degraded is True
|
||||
assert result.dense.timing.state is CandidateBackendState.AVAILABLE
|
||||
assert result.sparse.timing.state is CandidateBackendState.DEGRADED
|
||||
assert result.sparse.timing.error_code == "backend_unavailable"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_out_of_scope_backend_result_is_explicitly_degraded() -> None:
|
||||
result = await CandidateSearchService(
|
||||
RecordingBackend((candidate(CandidateBackendKind.DENSE, space_id=uuid4()),)),
|
||||
RecordingBackend(),
|
||||
).search(request())
|
||||
|
||||
assert result.dense.timing.state is CandidateBackendState.DEGRADED
|
||||
assert result.dense.timing.error_code == "candidate_contract_violation"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_backend_timeout_is_explicit_and_bounded() -> None:
|
||||
result = await CandidateSearchService(SlowBackend(), RecordingBackend()).search(
|
||||
request(timeout_seconds=0.05)
|
||||
)
|
||||
|
||||
assert result.degraded is True
|
||||
assert result.dense.timing.state is CandidateBackendState.TIMEOUT
|
||||
assert result.dense.timing.error_code == "backend_timeout"
|
||||
|
||||
|
||||
def test_candidate_limits_are_enforced_before_backend_execution() -> None:
|
||||
with pytest.raises(ValueError):
|
||||
request(dense_limit=51)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_qdrant_backend_decodes_provenance_and_uses_named_active_alias() -> None:
|
||||
point_id = uuid4()
|
||||
workspace_id = uuid4()
|
||||
space_id = uuid4()
|
||||
provenance = {
|
||||
"workspace_id": str(workspace_id),
|
||||
"space_id": str(space_id),
|
||||
"document_id": str(uuid4()),
|
||||
"document_version_id": str(uuid4()),
|
||||
"chunk_id": str(uuid4()),
|
||||
"source_id": str(uuid4()),
|
||||
"parent_chunk_id": None,
|
||||
"status": "published",
|
||||
"content_type": "text/plain",
|
||||
"language": "en",
|
||||
"trust_level": "source",
|
||||
"embedding_model_digest": PROFILE.model_digest,
|
||||
"chunk_profile_version": PROFILE.version,
|
||||
}
|
||||
|
||||
class StaticQdrantClient:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[dict[str, object]] = []
|
||||
|
||||
async def get_aliases(self, **kwargs: object) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
aliases=[
|
||||
SimpleNamespace(
|
||||
alias_name=PROFILE.alias_name,
|
||||
collection_name=PROFILE.collection_name,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
async def get_collection(self, collection_name: str, **kwargs: object) -> SimpleNamespace:
|
||||
assert collection_name == PROFILE.collection_name
|
||||
return SimpleNamespace(
|
||||
config=SimpleNamespace(
|
||||
params=SimpleNamespace(
|
||||
vectors={
|
||||
"dense": models.VectorParams(
|
||||
size=PROFILE.dimensions,
|
||||
distance=models.Distance.COSINE,
|
||||
)
|
||||
},
|
||||
sparse_vectors={"sparse": models.SparseVectorParams()},
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
async def query_points(self, **kwargs: object) -> models.QueryResponse:
|
||||
self.calls.append(kwargs)
|
||||
return models.QueryResponse(
|
||||
points=[
|
||||
models.ScoredPoint(id=str(point_id), version=0, score=0.9, payload=provenance)
|
||||
]
|
||||
)
|
||||
|
||||
client = StaticQdrantClient()
|
||||
backend = QdrantCandidateBackend(client, CandidateBackendKind.DENSE)
|
||||
plan = compile_retrieval_filter(AuthorizedRetrievalScope(workspace_id, frozenset({space_id})))
|
||||
|
||||
results = await backend.search(
|
||||
BackendSearchRequest(
|
||||
backend=CandidateBackendKind.DENSE,
|
||||
profile=PROFILE,
|
||||
query=DenseQuery((0.1, 0.2)),
|
||||
filter_plan=plan,
|
||||
limit=1,
|
||||
)
|
||||
)
|
||||
|
||||
assert results[0].point_id == point_id
|
||||
assert str(results[0].provenance.document_id) == provenance["document_id"]
|
||||
assert client.calls[0]["collection_name"] == PROFILE.collection_name
|
||||
assert client.calls[0]["using"] == "dense"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_internal_shadow_target_reads_named_collection_without_alias_cutover() -> None:
|
||||
shadow_space = uuid4()
|
||||
shadow_profile = CollectionProfile(
|
||||
kind=CollectionKind.DENSE,
|
||||
profile="modelforge-rag",
|
||||
version="m6",
|
||||
dimensions=2,
|
||||
model_digest="a" * 64,
|
||||
embedding_space_id=shadow_space,
|
||||
)
|
||||
workspace_id, space_id, point_id = uuid4(), uuid4(), uuid4()
|
||||
|
||||
class ShadowClient:
|
||||
async def get_aliases(self) -> SimpleNamespace:
|
||||
return SimpleNamespace(aliases=[])
|
||||
|
||||
async def get_collection(self, collection_name: str) -> SimpleNamespace:
|
||||
assert collection_name == shadow_profile.collection_name
|
||||
return SimpleNamespace(
|
||||
config=SimpleNamespace(
|
||||
params=SimpleNamespace(
|
||||
vectors={
|
||||
"dense": models.VectorParams(size=2, distance=models.Distance.COSINE)
|
||||
},
|
||||
sparse_vectors={"sparse": models.SparseVectorParams()},
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
async def query_points(self, **kwargs: object) -> models.QueryResponse:
|
||||
return models.QueryResponse(
|
||||
points=[
|
||||
models.ScoredPoint(
|
||||
id=str(point_id),
|
||||
version=0,
|
||||
score=0.9,
|
||||
payload={
|
||||
"workspace_id": str(workspace_id),
|
||||
"space_id": str(space_id),
|
||||
"document_id": str(uuid4()),
|
||||
"document_version_id": str(uuid4()),
|
||||
"chunk_id": str(point_id),
|
||||
"source_id": str(uuid4()),
|
||||
"parent_chunk_id": None,
|
||||
"status": "published",
|
||||
"content_type": "text/plain",
|
||||
"language": "en",
|
||||
"trust_level": "source",
|
||||
"embedding_model_digest": shadow_profile.model_digest,
|
||||
"chunk_profile_version": shadow_profile.version,
|
||||
"embedding_space_id": str(shadow_space),
|
||||
},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
backend = QdrantCandidateBackend(
|
||||
ShadowClient(),
|
||||
CandidateBackendKind.DENSE,
|
||||
target=EvaluationRetrievalTarget.SHADOW,
|
||||
)
|
||||
results = await backend.search(
|
||||
BackendSearchRequest(
|
||||
backend=CandidateBackendKind.DENSE,
|
||||
profile=shadow_profile,
|
||||
query=DenseQuery((1.0, 0.0)),
|
||||
filter_plan=compile_retrieval_filter(
|
||||
AuthorizedRetrievalScope(workspace_id, frozenset({space_id}))
|
||||
),
|
||||
limit=1,
|
||||
)
|
||||
)
|
||||
assert results[0].point_id == point_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_current_target_reads_the_explicit_pinned_collection_and_alias() -> None:
|
||||
pinned_collection = "rag_dense_current_pinned_20260912"
|
||||
point_id = uuid4()
|
||||
workspace_id, space_id = uuid4(), uuid4()
|
||||
|
||||
class PinnedClient:
|
||||
def __init__(self) -> None:
|
||||
self.collection_names: list[str] = []
|
||||
|
||||
async def get_aliases(self) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
aliases=[
|
||||
SimpleNamespace(
|
||||
alias_name=PROFILE.alias_name,
|
||||
collection_name=pinned_collection,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
async def get_collection(self, collection_name: str) -> SimpleNamespace:
|
||||
self.collection_names.append(collection_name)
|
||||
assert collection_name == pinned_collection
|
||||
return SimpleNamespace(
|
||||
config=SimpleNamespace(
|
||||
params=SimpleNamespace(
|
||||
vectors={
|
||||
"dense": models.VectorParams(
|
||||
size=PROFILE.dimensions, distance=models.Distance.COSINE
|
||||
)
|
||||
},
|
||||
sparse_vectors={"sparse": models.SparseVectorParams()},
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
async def query_points(self, **kwargs: object) -> models.QueryResponse:
|
||||
assert kwargs["collection_name"] == pinned_collection
|
||||
return models.QueryResponse(
|
||||
points=[
|
||||
models.ScoredPoint(
|
||||
id=str(point_id),
|
||||
version=0,
|
||||
score=0.9,
|
||||
payload={
|
||||
"workspace_id": str(workspace_id),
|
||||
"space_id": str(space_id),
|
||||
"document_id": str(uuid4()),
|
||||
"document_version_id": str(uuid4()),
|
||||
"chunk_id": str(point_id),
|
||||
"source_id": str(uuid4()),
|
||||
"parent_chunk_id": None,
|
||||
"status": "published",
|
||||
"content_type": "text/plain",
|
||||
"language": "en",
|
||||
"trust_level": "source",
|
||||
"embedding_model_digest": PROFILE.model_digest,
|
||||
"chunk_profile_version": PROFILE.version,
|
||||
},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
client = PinnedClient()
|
||||
backend = QdrantCandidateBackend(
|
||||
client, CandidateBackendKind.DENSE, collection_name=pinned_collection
|
||||
)
|
||||
result = await backend.search(
|
||||
BackendSearchRequest(
|
||||
backend=CandidateBackendKind.DENSE,
|
||||
profile=PROFILE,
|
||||
query=DenseQuery((0.1, 0.2)),
|
||||
filter_plan=compile_retrieval_filter(
|
||||
AuthorizedRetrievalScope(workspace_id, frozenset({space_id}))
|
||||
),
|
||||
limit=1,
|
||||
)
|
||||
)
|
||||
|
||||
assert result[0].point_id == point_id
|
||||
assert client.collection_names == [pinned_collection]
|
||||
|
||||
|
||||
def test_qdrant_backend_rejects_unsafe_pinned_collection_name() -> None:
|
||||
with pytest.raises(ValueError, match="safe Qdrant collection"):
|
||||
QdrantCandidateBackend(
|
||||
object(), CandidateBackendKind.DENSE, collection_name="../production"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dense_and_sparse_tasks_start_before_either_completes() -> None:
|
||||
started = [0]
|
||||
both_started = asyncio.Event()
|
||||
|
||||
class GatedBackend:
|
||||
async def search(self, request: BackendSearchRequest) -> tuple[Candidate, ...]:
|
||||
started[0] += 1
|
||||
if started[0] == 2:
|
||||
both_started.set()
|
||||
await both_started.wait()
|
||||
return ()
|
||||
|
||||
result = await CandidateSearchService(GatedBackend(), GatedBackend()).search(
|
||||
request(timeout_seconds=0.5)
|
||||
)
|
||||
|
||||
assert started[0] == 2
|
||||
assert result.degraded is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancellation_is_propagated_to_both_backends() -> None:
|
||||
started = asyncio.Event()
|
||||
cancelled = [0]
|
||||
blocked = asyncio.Event()
|
||||
|
||||
class CancellableBackend:
|
||||
async def search(self, request: BackendSearchRequest) -> tuple[Candidate, ...]:
|
||||
started.set()
|
||||
try:
|
||||
await blocked.wait()
|
||||
except asyncio.CancelledError:
|
||||
cancelled[0] += 1
|
||||
raise
|
||||
return ()
|
||||
|
||||
task = asyncio.create_task(
|
||||
CandidateSearchService(CancellableBackend(), CancellableBackend()).search(request())
|
||||
)
|
||||
await started.wait()
|
||||
task.cancel()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert cancelled[0] == 2
|
||||
@@ -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