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,352 @@
|
||||
"""ParentExpansionService: authorized, budget-bounded, dedup-safe parent attachment."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
from ragcore.application.retrieval.authorization import AuthorizedRetrievalScope
|
||||
from ragcore.domain.citations import Citation
|
||||
from ragcore.retrieval.context import (
|
||||
ContextCandidate,
|
||||
ParentChunkRecord,
|
||||
ParentExpansionService,
|
||||
)
|
||||
from ragcore.retrieval.context.models import ContextContractError
|
||||
|
||||
|
||||
def make_citation(
|
||||
workspace_id,
|
||||
space_id,
|
||||
document_id,
|
||||
version_id,
|
||||
chunk_id,
|
||||
run_id,
|
||||
excerpt,
|
||||
*,
|
||||
parent_chunk_id=None,
|
||||
):
|
||||
return Citation.create(
|
||||
workspace_id=workspace_id,
|
||||
space_id=space_id,
|
||||
document_id=document_id,
|
||||
document_version_id=version_id,
|
||||
chunk_id=chunk_id,
|
||||
parent_chunk_id=parent_chunk_id,
|
||||
location={"page_start": 1},
|
||||
excerpt=excerpt,
|
||||
source_uri="https://example.test/context.md",
|
||||
acquired_at=datetime(2026, 8, 1, 12, tzinfo=UTC),
|
||||
retrieval_run_id=run_id,
|
||||
index_alias="rag_dense_hybrid_active",
|
||||
profile_version="hybrid-v1",
|
||||
)
|
||||
|
||||
|
||||
def make_record(
|
||||
document_id,
|
||||
version_id,
|
||||
space_id,
|
||||
chunk_id,
|
||||
text,
|
||||
) -> ParentChunkRecord:
|
||||
return ParentChunkRecord(
|
||||
document_id=document_id,
|
||||
document_version_id=version_id,
|
||||
space_id=space_id,
|
||||
chunk_id=chunk_id,
|
||||
location={"page_start": 1},
|
||||
text=text,
|
||||
source_uri="https://example.test/parent.md",
|
||||
acquired_at=datetime(2026, 8, 1, 12, tzinfo=UTC),
|
||||
)
|
||||
|
||||
|
||||
class FakeLookup:
|
||||
"""Deterministic authorized-lookup fixture: records keyed by chunk id."""
|
||||
|
||||
def __init__(self, records: dict[UUID, ParentChunkRecord | None]) -> None:
|
||||
self._records = records
|
||||
self.calls: list[UUID] = []
|
||||
|
||||
async def get(self, parent_chunk_id, *, scope):
|
||||
self.calls.append(parent_chunk_id)
|
||||
return self._records.get(parent_chunk_id)
|
||||
|
||||
|
||||
def _scope(workspace_id, space_id) -> AuthorizedRetrievalScope:
|
||||
return AuthorizedRetrievalScope(workspace_id, frozenset({space_id}))
|
||||
|
||||
|
||||
async def test_authorized_parent_is_attached_and_improves_context() -> None:
|
||||
workspace_id, space_id, run_id = uuid4(), uuid4(), uuid4()
|
||||
document_id, version_id = uuid4(), uuid4()
|
||||
parent_id, child_id = uuid4(), uuid4()
|
||||
child_citation = make_citation(
|
||||
workspace_id, space_id, document_id, version_id, child_id, run_id,
|
||||
"child text", parent_chunk_id=parent_id,
|
||||
)
|
||||
candidate = ContextCandidate(child_citation, "child text", 1, 1.0)
|
||||
lookup = FakeLookup(
|
||||
{parent_id: make_record(document_id, version_id, space_id, parent_id, "parent text")}
|
||||
)
|
||||
service = ParentExpansionService(lookup)
|
||||
|
||||
(result,) = await service.expand(
|
||||
(candidate,),
|
||||
scope=_scope(workspace_id, space_id),
|
||||
retrieval_run_id=run_id,
|
||||
index_alias="rag_dense_hybrid_active",
|
||||
profile_version="hybrid-v1",
|
||||
)
|
||||
|
||||
assert result.parent_text == "parent text"
|
||||
assert result.parent_citation is not None
|
||||
assert result.parent_citation.chunk_id == parent_id
|
||||
assert result.parent_citation.excerpt == "parent text"
|
||||
|
||||
|
||||
async def test_missing_or_forbidden_parent_is_never_included() -> None:
|
||||
workspace_id, space_id, run_id = uuid4(), uuid4(), uuid4()
|
||||
document_id, version_id = uuid4(), uuid4()
|
||||
parent_id, child_id = uuid4(), uuid4()
|
||||
child_citation = make_citation(
|
||||
workspace_id, space_id, document_id, version_id, child_id, run_id,
|
||||
"child text", parent_chunk_id=parent_id,
|
||||
)
|
||||
candidate = ContextCandidate(child_citation, "child text", 1, 1.0)
|
||||
# The lookup port itself denies the parent (inactive/forbidden/not found).
|
||||
lookup = FakeLookup({parent_id: None})
|
||||
service = ParentExpansionService(lookup)
|
||||
|
||||
(result,) = await service.expand(
|
||||
(candidate,),
|
||||
scope=_scope(workspace_id, space_id),
|
||||
retrieval_run_id=run_id,
|
||||
index_alias="rag_dense_hybrid_active",
|
||||
profile_version="hybrid-v1",
|
||||
)
|
||||
|
||||
assert result.parent_text is None
|
||||
assert result.parent_citation is None
|
||||
|
||||
|
||||
async def test_parent_outside_effective_scope_is_rejected_defensively() -> None:
|
||||
workspace_id, space_id, forbidden_space, run_id = uuid4(), uuid4(), uuid4(), uuid4()
|
||||
document_id, version_id = uuid4(), uuid4()
|
||||
parent_id, child_id = uuid4(), uuid4()
|
||||
child_citation = make_citation(
|
||||
workspace_id, space_id, document_id, version_id, child_id, run_id,
|
||||
"child text", parent_chunk_id=parent_id,
|
||||
)
|
||||
candidate = ContextCandidate(child_citation, "child text", 1, 1.0)
|
||||
# A misbehaving/compromised lookup port claims a parent from a space the
|
||||
# caller never granted; the service must reject it regardless.
|
||||
lookup = FakeLookup(
|
||||
{parent_id: make_record(document_id, version_id, forbidden_space, parent_id, "leaked text")}
|
||||
)
|
||||
service = ParentExpansionService(lookup)
|
||||
|
||||
(result,) = await service.expand(
|
||||
(candidate,),
|
||||
scope=_scope(workspace_id, space_id),
|
||||
retrieval_run_id=run_id,
|
||||
index_alias="rag_dense_hybrid_active",
|
||||
profile_version="hybrid-v1",
|
||||
)
|
||||
|
||||
assert result.parent_citation is None
|
||||
assert result.parent_text is None
|
||||
|
||||
|
||||
async def test_parent_identical_to_child_text_is_not_attached() -> None:
|
||||
workspace_id, space_id, run_id = uuid4(), uuid4(), uuid4()
|
||||
document_id, version_id = uuid4(), uuid4()
|
||||
parent_id, child_id = uuid4(), uuid4()
|
||||
child_citation = make_citation(
|
||||
workspace_id, space_id, document_id, version_id, child_id, run_id,
|
||||
"same text", parent_chunk_id=parent_id,
|
||||
)
|
||||
candidate = ContextCandidate(child_citation, "same text", 1, 1.0)
|
||||
lookup = FakeLookup(
|
||||
{parent_id: make_record(document_id, version_id, space_id, parent_id, "same text")}
|
||||
)
|
||||
service = ParentExpansionService(lookup)
|
||||
|
||||
(result,) = await service.expand(
|
||||
(candidate,),
|
||||
scope=_scope(workspace_id, space_id),
|
||||
retrieval_run_id=run_id,
|
||||
index_alias="rag_dense_hybrid_active",
|
||||
profile_version="hybrid-v1",
|
||||
)
|
||||
|
||||
assert result.parent_citation is None
|
||||
|
||||
|
||||
async def test_candidate_without_a_parent_chunk_never_triggers_a_lookup() -> None:
|
||||
workspace_id, space_id, run_id = uuid4(), uuid4(), uuid4()
|
||||
document_id, version_id, chunk_id = uuid4(), uuid4(), uuid4()
|
||||
child_citation = make_citation(
|
||||
workspace_id, space_id, document_id, version_id, chunk_id, run_id, "top-level text"
|
||||
)
|
||||
candidate = ContextCandidate(child_citation, "top-level text", 1, 1.0)
|
||||
lookup = FakeLookup({})
|
||||
service = ParentExpansionService(lookup)
|
||||
|
||||
(result,) = await service.expand(
|
||||
(candidate,),
|
||||
scope=_scope(workspace_id, space_id),
|
||||
retrieval_run_id=run_id,
|
||||
index_alias="rag_dense_hybrid_active",
|
||||
profile_version="hybrid-v1",
|
||||
)
|
||||
|
||||
assert result.parent_citation is None
|
||||
assert lookup.calls == []
|
||||
|
||||
|
||||
async def test_shared_parent_is_expanded_only_once_avoiding_duplicate_text() -> None:
|
||||
workspace_id, space_id, run_id = uuid4(), uuid4(), uuid4()
|
||||
document_id, version_id = uuid4(), uuid4()
|
||||
parent_id = uuid4()
|
||||
first_child, second_child = uuid4(), uuid4()
|
||||
candidates = (
|
||||
ContextCandidate(
|
||||
make_citation(
|
||||
workspace_id, space_id, document_id, version_id, first_child, run_id,
|
||||
"first child", parent_chunk_id=parent_id,
|
||||
),
|
||||
"first child",
|
||||
1,
|
||||
1.0,
|
||||
),
|
||||
ContextCandidate(
|
||||
make_citation(
|
||||
workspace_id, space_id, document_id, version_id, second_child, run_id,
|
||||
"second child", parent_chunk_id=parent_id,
|
||||
),
|
||||
"second child",
|
||||
2,
|
||||
0.5,
|
||||
),
|
||||
)
|
||||
lookup = FakeLookup(
|
||||
{parent_id: make_record(document_id, version_id, space_id, parent_id, "shared parent")}
|
||||
)
|
||||
service = ParentExpansionService(lookup)
|
||||
|
||||
first, second = await service.expand(
|
||||
candidates,
|
||||
scope=_scope(workspace_id, space_id),
|
||||
retrieval_run_id=run_id,
|
||||
index_alias="rag_dense_hybrid_active",
|
||||
profile_version="hybrid-v1",
|
||||
)
|
||||
|
||||
assert first.parent_text == "shared parent"
|
||||
assert second.parent_text is None
|
||||
assert lookup.calls == [parent_id] # looked up once, not twice
|
||||
|
||||
|
||||
async def test_expansion_budget_is_bounded_and_enforced_in_order() -> None:
|
||||
workspace_id, space_id, run_id = uuid4(), uuid4(), uuid4()
|
||||
document_id, version_id = uuid4(), uuid4()
|
||||
parent_ids = [uuid4(), uuid4(), uuid4()]
|
||||
candidates = tuple(
|
||||
ContextCandidate(
|
||||
make_citation(
|
||||
workspace_id, space_id, document_id, version_id, uuid4(), run_id,
|
||||
f"child {index}", parent_chunk_id=parent_ids[index],
|
||||
),
|
||||
f"child {index}",
|
||||
index + 1,
|
||||
1.0,
|
||||
)
|
||||
for index in range(3)
|
||||
)
|
||||
records = {
|
||||
parent_id: make_record(document_id, version_id, space_id, parent_id, "x" * 10)
|
||||
for parent_id in parent_ids
|
||||
}
|
||||
lookup = FakeLookup(records)
|
||||
# Each parent's text is 10 word-like bytes; a budget of 15 allows only one.
|
||||
service = ParentExpansionService(lookup, max_expansion_tokens=15)
|
||||
|
||||
results = await service.expand(
|
||||
candidates,
|
||||
scope=_scope(workspace_id, space_id),
|
||||
retrieval_run_id=run_id,
|
||||
index_alias="rag_dense_hybrid_active",
|
||||
profile_version="hybrid-v1",
|
||||
)
|
||||
|
||||
expanded_count = sum(1 for item in results if item.parent_citation is not None)
|
||||
assert expanded_count == 1
|
||||
assert results[0].parent_citation is not None # rank-1 candidate expanded first
|
||||
|
||||
|
||||
async def test_expansion_is_deterministic_given_the_same_input() -> None:
|
||||
workspace_id, space_id, run_id = uuid4(), uuid4(), uuid4()
|
||||
document_id, version_id = uuid4(), uuid4()
|
||||
parent_id, child_id = uuid4(), uuid4()
|
||||
child_citation = make_citation(
|
||||
workspace_id, space_id, document_id, version_id, child_id, run_id,
|
||||
"child text", parent_chunk_id=parent_id,
|
||||
)
|
||||
candidate = ContextCandidate(child_citation, "child text", 1, 1.0)
|
||||
record = make_record(document_id, version_id, space_id, parent_id, "parent text")
|
||||
service_a = ParentExpansionService(FakeLookup({parent_id: record}))
|
||||
service_b = ParentExpansionService(FakeLookup({parent_id: record}))
|
||||
|
||||
(result_a,) = await service_a.expand(
|
||||
(candidate,),
|
||||
scope=_scope(workspace_id, space_id),
|
||||
retrieval_run_id=run_id,
|
||||
index_alias="rag_dense_hybrid_active",
|
||||
profile_version="hybrid-v1",
|
||||
)
|
||||
(result_b,) = await service_b.expand(
|
||||
(candidate,),
|
||||
scope=_scope(workspace_id, space_id),
|
||||
retrieval_run_id=run_id,
|
||||
index_alias="rag_dense_hybrid_active",
|
||||
profile_version="hybrid-v1",
|
||||
)
|
||||
|
||||
assert result_a == result_b
|
||||
|
||||
|
||||
async def test_max_expansion_tokens_is_bounded() -> None:
|
||||
with pytest.raises(ValueError, match="max_expansion_tokens"):
|
||||
ParentExpansionService(FakeLookup({}), max_expansion_tokens=0)
|
||||
|
||||
|
||||
async def test_invalid_token_counter_result_raises() -> None:
|
||||
workspace_id, space_id, run_id = uuid4(), uuid4(), uuid4()
|
||||
document_id, version_id = uuid4(), uuid4()
|
||||
parent_id, child_id = uuid4(), uuid4()
|
||||
child_citation = make_citation(
|
||||
workspace_id, space_id, document_id, version_id, child_id, run_id,
|
||||
"child text", parent_chunk_id=parent_id,
|
||||
)
|
||||
candidate = ContextCandidate(child_citation, "child text", 1, 1.0)
|
||||
lookup = FakeLookup(
|
||||
{parent_id: make_record(document_id, version_id, space_id, parent_id, "parent text")}
|
||||
)
|
||||
|
||||
class BrokenCounter:
|
||||
def count(self, text: str) -> int:
|
||||
return -1
|
||||
|
||||
service = ParentExpansionService(lookup, token_counter=BrokenCounter())
|
||||
with pytest.raises(ContextContractError, match="token counter"):
|
||||
await service.expand(
|
||||
(candidate,),
|
||||
scope=_scope(workspace_id, space_id),
|
||||
retrieval_run_id=run_id,
|
||||
index_alias="rag_dense_hybrid_active",
|
||||
profile_version="hybrid-v1",
|
||||
)
|
||||
Reference in new issue
Block a user