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,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
|
||||
Reference in new issue
Block a user