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,376 @@
|
||||
"""Fixtures for the MCP contract tests.
|
||||
|
||||
Mirrors the same fixture-construction pattern as `tests/contract/api/**`
|
||||
and `tests/contract/sdk/python/**` (real `RetrievalAuthorizationService`,
|
||||
a fixture executor/generator standing in only for the retrieval
|
||||
backend/model provider) -- these tests exist to prove the MCP adapter's
|
||||
own auth/mapping/rate-limit layer, not to re-prove authorization/
|
||||
retrieval business logic already covered there.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass, replace
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
import httpx2
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.client.streamable_http import streamable_http_client
|
||||
from mcp.server.mcpserver import MCPServer
|
||||
|
||||
from ragcore.api.v1.answers.service import AnswerService
|
||||
from ragcore.api.v1.context.service import (
|
||||
ContextExecutionRequest,
|
||||
ContextExecutionResult,
|
||||
ContextService,
|
||||
)
|
||||
from ragcore.api.v1.search.service import (
|
||||
AuthorizationInputs,
|
||||
SearchHit,
|
||||
SearchService,
|
||||
)
|
||||
from ragcore.application.identity.credentials import ServiceAccountCredentialService
|
||||
from ragcore.application.retrieval.authorization import RetrievalAuthorizationService
|
||||
from ragcore.domain.authorization import Principal
|
||||
from ragcore.domain.citations import Citation
|
||||
from ragcore.domain.control import Application, Role, Scope, ServiceAccount, Workspace
|
||||
from ragcore.infrastructure.database.identity import InMemoryCredentialRepository
|
||||
from ragcore.mcp.server import MCPDependencies, build_mcp_server
|
||||
from ragcore.retrieval.context import ContextCandidate
|
||||
from ragcore.retrieval.ranking import FusedCandidate
|
||||
from ragcore.retrieval.search.models import Candidate, CandidateBackendKind, CandidateProvenance
|
||||
|
||||
EVIDENCE_EXCERPT = "The refund window is thirty days from the original purchase date."
|
||||
CREDENTIAL_PEPPER = "test-credential-pepper-that-is-long-enough"
|
||||
|
||||
|
||||
class InMemoryDocumentRepository:
|
||||
def __init__(self) -> None:
|
||||
self._documents: dict[UUID, Any] = {}
|
||||
self._versions: dict[UUID, Any] = {}
|
||||
|
||||
def add_document(self, document: Any) -> None:
|
||||
self._documents[document.id] = document
|
||||
|
||||
def add_version(self, version: Any) -> None:
|
||||
self._versions[version.id] = version
|
||||
|
||||
async def get_document(self, document_id: UUID) -> Any:
|
||||
return self._documents.get(document_id)
|
||||
|
||||
async def get_version(self, version_id: UUID) -> Any:
|
||||
return self._versions.get(version_id)
|
||||
|
||||
async def list_searchable_versions(self, document_id: UUID) -> tuple[Any, ...]:
|
||||
from ragcore.application.documents.models import DocumentVersionState
|
||||
|
||||
return tuple(
|
||||
v
|
||||
for v in self._versions.values()
|
||||
if v.document_id == document_id and v.state is DocumentVersionState.PUBLISHED
|
||||
)
|
||||
|
||||
|
||||
class InMemoryCitationRepository:
|
||||
def __init__(self) -> None:
|
||||
self._citations: dict[UUID, Citation] = {}
|
||||
|
||||
def add(self, citation: Citation) -> None:
|
||||
self._citations[citation.id] = citation
|
||||
|
||||
async def get_authorized(
|
||||
self, citation_id: UUID, *, workspace_id: UUID, effective_space_ids: frozenset[UUID]
|
||||
) -> Citation | None:
|
||||
citation = self._citations.get(citation_id)
|
||||
if citation is None or citation.workspace_id != workspace_id:
|
||||
return None
|
||||
if citation.space_id not in effective_space_ids:
|
||||
return None
|
||||
return citation
|
||||
|
||||
|
||||
class _FixtureInputs:
|
||||
"""Loads the one space plus a real workspace-level `Grant` tying it to
|
||||
`application` -- `EffectiveSpaceResolver` requires an actual `Grant`
|
||||
row (not just `Application.default_knowledge_space_ids`, which is only
|
||||
consulted when the caller's own `requested_space_ids` is empty) to
|
||||
intersect an *explicit* request against for a SERVICE_ACCOUNT
|
||||
principal, exactly like a real workspace-level grant REST would use.
|
||||
"""
|
||||
|
||||
def __init__(self, workspace: Workspace, space: Any, application: Application) -> None:
|
||||
self.workspace = workspace
|
||||
self.space = space
|
||||
self.application = application
|
||||
self.grant = workspace.grant(
|
||||
id=uuid4(), application=application, knowledge_space=space, role=Role.EDITOR
|
||||
)
|
||||
|
||||
async def load(self, principal: Principal) -> Any:
|
||||
return type(
|
||||
"Inputs",
|
||||
(),
|
||||
{
|
||||
"spaces": (self.space,),
|
||||
"grants": (self.grant,),
|
||||
"user_grants": (),
|
||||
"application": self.application,
|
||||
},
|
||||
)()
|
||||
|
||||
|
||||
@dataclass
|
||||
class McpFixture:
|
||||
server: MCPServer
|
||||
deps: MCPDependencies
|
||||
workspace: Workspace
|
||||
application: Application
|
||||
space_id: UUID
|
||||
token: str
|
||||
document_repository: InMemoryDocumentRepository
|
||||
citation_repository: InMemoryCitationRepository
|
||||
|
||||
|
||||
def _citation(workspace_id: UUID, space_id: UUID) -> Citation:
|
||||
return Citation.create(
|
||||
workspace_id=workspace_id,
|
||||
space_id=space_id,
|
||||
document_id=uuid4(),
|
||||
document_version_id=uuid4(),
|
||||
chunk_id=uuid4(),
|
||||
parent_chunk_id=None,
|
||||
location={"page_start": 4, "section": "Refunds"},
|
||||
excerpt=EVIDENCE_EXCERPT,
|
||||
source_uri="https://example.test/policy.md",
|
||||
acquired_at=datetime(2026, 8, 1, 12, tzinfo=UTC),
|
||||
retrieval_run_id=uuid4(),
|
||||
index_alias="rag_dense_hybrid_active",
|
||||
profile_version="hybrid-v1",
|
||||
)
|
||||
|
||||
|
||||
async def build_mcp_fixture(*, grant_scopes: set[Scope] | None = None) -> McpFixture:
|
||||
workspace = Workspace(uuid4(), "MCP workspace", "mcp-workspace")
|
||||
space_id = uuid4()
|
||||
space = workspace.create_knowledge_space(id=space_id, name="MCP space", slug="mcp-space")
|
||||
application = workspace.create_application(
|
||||
id=uuid4(),
|
||||
name="MCP client",
|
||||
slug="mcp-client",
|
||||
scopes=frozenset(
|
||||
grant_scopes
|
||||
or {
|
||||
Scope.SEARCH,
|
||||
Scope.CONTEXT,
|
||||
Scope.ANSWER,
|
||||
Scope.DOCUMENTS_READ,
|
||||
Scope.CITATIONS_READ,
|
||||
}
|
||||
),
|
||||
default_knowledge_space_ids=frozenset({space_id}),
|
||||
)
|
||||
service_account = ServiceAccount(
|
||||
id=uuid4(),
|
||||
workspace_id=workspace.id,
|
||||
application_id=application.id,
|
||||
name="MCP service account",
|
||||
scopes=application.scopes,
|
||||
)
|
||||
|
||||
async def load_service_account(account_id: UUID) -> ServiceAccount | None:
|
||||
return service_account if account_id == service_account.id else None
|
||||
|
||||
credential_service = ServiceAccountCredentialService(
|
||||
InMemoryCredentialRepository(),
|
||||
CREDENTIAL_PEPPER,
|
||||
service_account_loader=load_service_account,
|
||||
)
|
||||
issued = await credential_service.issue(service_account, scopes=service_account.scopes)
|
||||
|
||||
citation = _citation(workspace.id, space_id)
|
||||
candidate = Candidate(
|
||||
point_id=uuid4(),
|
||||
score=0.87,
|
||||
rank=1,
|
||||
backend=CandidateBackendKind.DENSE,
|
||||
provenance=CandidateProvenance(
|
||||
workspace_id=workspace.id,
|
||||
space_id=space_id,
|
||||
document_id=citation.document_id,
|
||||
document_version_id=citation.document_version_id,
|
||||
chunk_id=citation.chunk_id,
|
||||
source_id=uuid4(),
|
||||
parent_chunk_id=None,
|
||||
status="published",
|
||||
content_type="text/markdown",
|
||||
language="en",
|
||||
trust_level="trusted_internal",
|
||||
),
|
||||
)
|
||||
hit = SearchHit(
|
||||
fused=FusedCandidate(candidate, 1 / 61, 1, None),
|
||||
text=EVIDENCE_EXCERPT,
|
||||
title="Refund policy",
|
||||
citation=citation,
|
||||
dense_score=0.87,
|
||||
)
|
||||
|
||||
class _SearchExecutor:
|
||||
async def execute(self, request: Any) -> Any:
|
||||
refreshed = replace(
|
||||
hit,
|
||||
citation=Citation.create(
|
||||
workspace_id=hit.citation.workspace_id,
|
||||
space_id=hit.citation.space_id,
|
||||
document_id=hit.citation.document_id,
|
||||
document_version_id=hit.citation.document_version_id,
|
||||
chunk_id=hit.citation.chunk_id,
|
||||
parent_chunk_id=hit.citation.parent_chunk_id,
|
||||
location=hit.citation.location,
|
||||
excerpt=hit.citation.excerpt,
|
||||
source_uri=hit.citation.source_uri,
|
||||
acquired_at=hit.citation.acquired_at,
|
||||
retrieval_run_id=request.retrieval_run_id,
|
||||
index_alias=hit.citation.index_alias,
|
||||
profile_version=hit.citation.profile_version,
|
||||
),
|
||||
)
|
||||
return type(
|
||||
"Result",
|
||||
(),
|
||||
{
|
||||
"hits": (refreshed,),
|
||||
"degraded": False,
|
||||
"degradation_codes": (),
|
||||
"timings_ms": {},
|
||||
},
|
||||
)()
|
||||
|
||||
inputs = _FixtureInputs(workspace, space, application)
|
||||
search_service = SearchService(
|
||||
authorization=RetrievalAuthorizationService(),
|
||||
inputs=inputs,
|
||||
executor=_SearchExecutor(),
|
||||
)
|
||||
|
||||
context_candidate = ContextCandidate(citation, EVIDENCE_EXCERPT, 1, 1.0)
|
||||
|
||||
class _ContextExecutor:
|
||||
async def execute(self, request: ContextExecutionRequest) -> ContextExecutionResult:
|
||||
refreshed = Citation.create(
|
||||
workspace_id=citation.workspace_id,
|
||||
space_id=citation.space_id,
|
||||
document_id=citation.document_id,
|
||||
document_version_id=citation.document_version_id,
|
||||
chunk_id=citation.chunk_id,
|
||||
parent_chunk_id=citation.parent_chunk_id,
|
||||
location=citation.location,
|
||||
excerpt=citation.excerpt,
|
||||
source_uri=citation.source_uri,
|
||||
acquired_at=citation.acquired_at,
|
||||
retrieval_run_id=request.retrieval_run_id,
|
||||
index_alias=citation.index_alias,
|
||||
profile_version=citation.profile_version,
|
||||
)
|
||||
return ContextExecutionResult(
|
||||
candidates=(replace(context_candidate, citation=refreshed),)
|
||||
)
|
||||
|
||||
context_service = ContextService(
|
||||
authorization=RetrievalAuthorizationService(), inputs=inputs, executor=_ContextExecutor()
|
||||
)
|
||||
|
||||
import json
|
||||
import re
|
||||
|
||||
citation_id_pattern = re.compile(r"citation_ids=([0-9a-fA-F-]{36})")
|
||||
|
||||
class _Generator:
|
||||
async def generate(self, *, prompt: str, generation_profile: str, spaces: Any) -> str:
|
||||
match = citation_id_pattern.search(prompt)
|
||||
assert match is not None
|
||||
return json.dumps(
|
||||
{
|
||||
"status": "answerable",
|
||||
"answer": "The refund window is thirty days from purchase.",
|
||||
"citations": [match.group(1)],
|
||||
}
|
||||
)
|
||||
|
||||
answer_service = AnswerService(
|
||||
authorization=RetrievalAuthorizationService(),
|
||||
inputs=inputs,
|
||||
context_executor=_ContextExecutor(),
|
||||
generator=_Generator(),
|
||||
)
|
||||
|
||||
document_repository = InMemoryDocumentRepository()
|
||||
citation_repository = InMemoryCitationRepository()
|
||||
citation_repository.add(citation)
|
||||
|
||||
class _AuthorizationInputsForMcp:
|
||||
async def load(self, principal: Principal) -> AuthorizationInputs:
|
||||
loaded = await inputs.load(principal)
|
||||
return AuthorizationInputs(
|
||||
spaces=loaded.spaces,
|
||||
grants=loaded.grants,
|
||||
user_grants=loaded.user_grants,
|
||||
application=loaded.application,
|
||||
)
|
||||
|
||||
deps = MCPDependencies(
|
||||
search_application=search_service,
|
||||
context_application=context_service,
|
||||
answer_application=answer_service,
|
||||
authorization_inputs=_AuthorizationInputsForMcp(),
|
||||
credential_service=credential_service,
|
||||
citation_repository=citation_repository,
|
||||
document_repository=document_repository,
|
||||
)
|
||||
server = build_mcp_server(deps)
|
||||
|
||||
return McpFixture(
|
||||
server=server,
|
||||
deps=deps,
|
||||
workspace=workspace,
|
||||
application=application,
|
||||
space_id=space_id,
|
||||
token=issued.token,
|
||||
document_repository=document_repository,
|
||||
citation_repository=citation_repository,
|
||||
)
|
||||
|
||||
|
||||
_MCP_TEST_URL = "http://127.0.0.1:8765"
|
||||
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def mcp_session(
|
||||
server: MCPServer, *, token: str | None = None
|
||||
) -> AsyncIterator[ClientSession]:
|
||||
"""A real streamable-HTTP MCP session over `httpx2.ASGITransport` -- genuine
|
||||
Starlette `Request`/header handling, no real TCP socket needed.
|
||||
|
||||
`ASGITransport` never runs the ASGI app's own lifespan protocol (that's
|
||||
normally what starts the session manager's task group when a real
|
||||
server like uvicorn owns the app), so `session_manager.run()` is
|
||||
entered explicitly here instead, mirroring what a real deployment's
|
||||
lifespan would do.
|
||||
"""
|
||||
|
||||
app = server.streamable_http_app(stateless_http=True)
|
||||
headers = {"authorization": f"Bearer {token}"} if token else {}
|
||||
async with (
|
||||
server.session_manager.run(),
|
||||
httpx2.AsyncClient(
|
||||
transport=httpx2.ASGITransport(app=app), base_url=_MCP_TEST_URL, headers=headers
|
||||
) as http_client,
|
||||
streamable_http_client(f"{_MCP_TEST_URL}/mcp", http_client=http_client) as (read, write),
|
||||
ClientSession(read, write) as session,
|
||||
):
|
||||
await session.initialize()
|
||||
yield session
|
||||
Reference in new issue
Block a user