diff --git a/backend/alembic/versions/0060_knowledge_base_rerank.py b/backend/alembic/versions/0060_knowledge_base_rerank.py new file mode 100644 index 000000000..cc4b9c3d8 --- /dev/null +++ b/backend/alembic/versions/0060_knowledge_base_rerank.py @@ -0,0 +1,61 @@ +"""Which reranker - and whose key - a collection reranks search results with. + +Retrieval fetches candidates by vector similarity and returns the top ones. A +reranker is a second pass: a model scores each candidate against the query +directly and reorders them, which is more accurate than the distance the vector +index sorts by. It is optional and off by default. + +Two nullable columns, mirroring the embedding pair. `rerank_model` is the +reranker's name; `rerank_secret_id` is the organization vault key that pays for +it. Reranking runs only when *both* are set - either NULL leaves retrieval +exactly as it was, so existing rows and unconfigured deployments are unchanged +by this migration. SET NULL on delete for the same reason the embedding key is: +losing the key drops reranking, it does not take document search down. Unlike +the embedding key there is no deployment fallback - a reranker with no key is +simply off. + +Revision ID: 0060_knowledge_base_rerank +Revises: 0059_invite_fk_ondelete +Create Date: 2026-08-18 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa + +from alembic import op + +revision: str = "0060_knowledge_base_rerank" +down_revision: str | None = "0059_invite_fk_ondelete" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + op.add_column( + "knowledge_bases", + sa.Column("rerank_model", sa.String(length=128), nullable=True), + ) + op.add_column( + "knowledge_bases", + sa.Column("rerank_secret_id", sa.UUID(), nullable=True), + ) + op.create_foreign_key( + op.f("knowledge_bases_rerank_secret_id_fkey"), + "knowledge_bases", + "organization_secrets", + ["rerank_secret_id"], + ["id"], + ondelete="SET NULL", + ) + + +def downgrade() -> None: + op.drop_constraint( + op.f("knowledge_bases_rerank_secret_id_fkey"), + "knowledge_bases", + type_="foreignkey", + ) + op.drop_column("knowledge_bases", "rerank_secret_id") + op.drop_column("knowledge_bases", "rerank_model") diff --git a/backend/alembic/versions/0061_ingestion_spend_source.py b/backend/alembic/versions/0061_ingestion_spend_source.py new file mode 100644 index 000000000..80e958c22 --- /dev/null +++ b/backend/alembic/versions/0061_ingestion_spend_source.py @@ -0,0 +1,43 @@ +"""Tag each non-run RAG spend row as indexing or retrieval. + +`ingestion_spend` began as indexing alone, then a metered `POST /rag/search` +landed its embedding and rerank cost in the same table - both are RAG spend +outside any agent run. Left undistinguished, a search inflated the dashboard's +"indexing" subtotal. `source` tells them apart; both still count toward the +monthly budget, only the reporting split reads the column. + +Every row that predates the column is indexing, so `server_default` backfills +them to `'ingestion'` without a data migration. + +Revision ID: 0061_ingestion_spend_source +Revises: 0060_knowledge_base_rerank +Create Date: 2026-08-20 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa + +from alembic import op + +revision: str = "0061_ingestion_spend_source" +down_revision: str | None = "0060_knowledge_base_rerank" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + op.add_column( + "ingestion_spend", + sa.Column( + "source", + sa.String(length=16), + nullable=False, + server_default="ingestion", + ), + ) + + +def downgrade() -> None: + op.drop_column("ingestion_spend", "source") diff --git a/backend/app/agents/capabilities/budget/__init__.py b/backend/app/agents/capabilities/budget/__init__.py index 20bb45cad..3e7a6c4c0 100644 --- a/backend/app/agents/capabilities/budget/__init__.py +++ b/backend/app/agents/capabilities/budget/__init__.py @@ -9,6 +9,7 @@ SpendLedger, SpendLimit, SpendShare, + book_ambient_spend, booked_to, metered_by, price_request, @@ -26,6 +27,7 @@ "SpendLedger", "SpendLimit", "SpendShare", + "book_ambient_spend", "booked_to", "metered_by", "price_request", diff --git a/backend/app/agents/capabilities/budget/_capability.py b/backend/app/agents/capabilities/budget/_capability.py index 9e4c4446d..16cbb2c40 100644 --- a/backend/app/agents/capabilities/budget/_capability.py +++ b/backend/app/agents/capabilities/budget/_capability.py @@ -434,6 +434,24 @@ def record_ambient_usage( ledger.record(model_name, usage, provider) +def book_ambient_spend(entry: SpendEntry) -> None: + """Book one already-priced entry to whichever ledger is active, if any is. + + The sibling of :func:`record_ambient_usage` for spend that is not + token-priced. `record` prices through `genai-prices`, which knows chat and + embedding models and nothing else, so a reranker call routed through it + would book `cost_usd=0, priced=False`. A reranker computes its own cost from + a published per-search price and hands the finished :class:`SpendEntry` here + instead, so the entry lands `priced=True` with a real number. + + A no-op when nothing is metering, for the same reason the sibling is: a + provider should not refuse to work because no ledger is open. + """ + ledger = _active_ledger.get() + if ledger is not None: + ledger.book(entry) + + def usage_counts(usage: RunUsage) -> tuple[int, int, int, int]: """The four counters a price is computed from, read off the run's usage. diff --git a/backend/app/agents/capabilities/knowledge/_search.py b/backend/app/agents/capabilities/knowledge/_search.py index a0ca6855e..76dd4bd2c 100644 --- a/backend/app/agents/capabilities/knowledge/_search.py +++ b/backend/app/agents/capabilities/knowledge/_search.py @@ -3,10 +3,12 @@ import contextvars import logging from typing import TYPE_CHECKING, Any +from uuid import UUID from app.core.config import settings from app.core.exceptions import AppException, ExternalServiceError from app.services.rag.embeddings import EmbeddingService +from app.services.rag.reranker import build_reranker from app.services.rag.retrieval import RetrievalService from app.services.rag.vectorstore import process_vector_store @@ -31,8 +33,14 @@ def get_retrieval_service() -> "BaseRetrievalService": rag_settings = settings.rag embedding_service = EmbeddingService(rag_settings) + # The reranker resolver is wired here too, not only on the /rag/search + # route: an agent's knowledge search reranks when its collection is + # configured, and the run's open ledger books the cost - which is the + # agent-run half of "spend recorded on both paths". _retrieval_service = RetrievalService( - process_vector_store(rag_settings, embedding_service), rag_settings + process_vector_store(rag_settings, embedding_service), + rag_settings, + reranker_resolver=build_reranker, ) return _retrieval_service @@ -87,6 +95,9 @@ async def search_knowledge_base( query: str, kb_collection_names: list[str] | None = None, top_k: int = 5, + *, + organization_id: UUID | None, + kb_collection_ids: list[UUID] | None = None, ) -> str: """Search the knowledge base and return formatted results. @@ -96,19 +107,40 @@ async def search_knowledge_base( agent's spec. Never supplied by the LLM directly - injected via PydanticAI Deps or the _active_kb_collections ContextVar. top_k: Number of top results to retrieve (default: 5). + organization_id: The organization the run acts for, so a collection name + shared across tenants resolves this one's embedding and rerank config + rather than another's (#913). + kb_collection_ids: The bound knowledge base id for each name, aligned by + index, so a name shared by another row in the same organization + resolves the bound row rather than whichever the name selects first + (#913). Absent - the ContextVar fallback, which carries no ids - each + collection falls back to the organization-scoped lookup. """ resolved = kb_collection_names if kb_collection_names else (_active_kb_collections.get() or []) if not resolved: return "No active knowledge bases selected for this conversation." + ids = kb_collection_ids or [] + aligned_ids = ids if len(ids) == len(resolved) else None + service: Any = get_retrieval_service() one_collection = len(resolved) == 1 try: if one_collection: - results = await service.retrieve(query=query, collection_name=resolved[0], limit=top_k) + results = await service.retrieve( + query=query, + collection_name=resolved[0], + limit=top_k, + organization_id=organization_id, + knowledge_base_id=aligned_ids[0] if aligned_ids else None, + ) else: results = await service.retrieve_multi( - query=query, collection_names=resolved, limit=top_k + query=query, + collection_names=resolved, + limit=top_k, + organization_id=organization_id, + knowledge_base_ids=aligned_ids, ) except AppException: # Already an account of what is wrong and what to do about it - an diff --git a/backend/app/agents/capabilities/knowledge/_toolset.py b/backend/app/agents/capabilities/knowledge/_toolset.py index 5b6ae1ee7..e1cd7b863 100644 --- a/backend/app/agents/capabilities/knowledge/_toolset.py +++ b/backend/app/agents/capabilities/knowledge/_toolset.py @@ -44,7 +44,14 @@ async def search_documents( # Resolved server-side from the agent's bound collections. The # model chooses *what* to search, never *where*. kb_collection_names=ctx.deps.kb_collection_names, + # The bound knowledge base ids, aligned with the names, so a + # shared collection name resolves the bound row's config and key + # rather than another same-named row's (#913). + kb_collection_ids=ctx.deps.kb_collection_ids, top_k=top_k or default_top_k, + # The run's own organization, so a collection name shared with + # another tenant resolves this agent's config, not theirs (#913). + organization_id=ctx.deps.organization_id, ) except Exception: # A retry rather than a returned message: an error in the shape of a diff --git a/backend/app/agents/capabilities/subagents/_capability.py b/backend/app/agents/capabilities/subagents/_capability.py index a7d04f339..f9be3180d 100644 --- a/backend/app/agents/capabilities/subagents/_capability.py +++ b/backend/app/agents/capabilities/subagents/_capability.py @@ -514,6 +514,7 @@ def _own_deps(self, kwargs: dict[str, Any]) -> dict[str, Any]: "deps": replace( clone, kb_collection_names=list(self._delegate.collection_names), + kb_collection_ids=list(self._delegate.collection_ids), request_approval=( None if self._journal.in_background() else clone.request_approval ), diff --git a/backend/app/agents/deps.py b/backend/app/agents/deps.py index f40efae6a..238a8dfe8 100644 --- a/backend/app/agents/deps.py +++ b/backend/app/agents/deps.py @@ -56,8 +56,13 @@ class AgentDeps: agent_id: UUID | None = None run_id: UUID | None = None - # Collection names this agent may search, resolved from its bindings. + # Collection names this agent may search, resolved from its bindings, and + # the id of the knowledge base each name was authorized as, aligned by + # index. `collection_name` is not unique, so search resolves each collection + # by its bound id rather than re-selecting by name (#913); the id is what + # keeps a shared name from resolving another row's config and key. kb_collection_names: list[str] = field(default_factory=list) + kb_collection_ids: list[UUID] = field(default_factory=list) # Set when the surface can ask the user something mid-run (WebSocket chat); # None on surfaces that cannot, so tools must handle its absence. @@ -109,10 +114,11 @@ def clone_for_subagent(self, max_depth: int = 0) -> AgentDeps: `subagent_events` - so a specialist's own delegation still narrates, one `depth` further in. - What it does **not** inherit: `kb_collection_names`. Those come from the - delegate's own spec, and inheriting the parent's would hand a specialist a - collection nobody granted it. The delegate's own are put back by the - delegation capability, from `ResolvedSubagent.collection_names` - because + What it does **not** inherit: `kb_collection_names` and their + `kb_collection_ids`. Those come from the delegate's own spec, and + inheriting the parent's would hand a specialist a collection nobody + granted it. The delegate's own are put back by the delegation + capability, from `ResolvedSubagent.collection_names`/`collection_ids` - because the deps our factory built for the child are *this* object's replacement, so the collections resolved for it would otherwise be resolved and never read. diff --git a/backend/app/agents/factory.py b/backend/app/agents/factory.py index d88dd8796..b3353f74a 100644 --- a/backend/app/agents/factory.py +++ b/backend/app/agents/factory.py @@ -235,6 +235,7 @@ def build_agent( # Read from `resources` rather than a parameter of its own: two sources # for one list is how they drift apart. kb_collection_names=list((resources or {}).get("kb_collection_names") or []), + kb_collection_ids=list((resources or {}).get("kb_collection_ids") or []), request_approval=request_approval, ) diff --git a/backend/app/agents/subagent_runtime.py b/backend/app/agents/subagent_runtime.py index dbc3976f5..5c7c301c1 100644 --- a/backend/app/agents/subagent_runtime.py +++ b/backend/app/agents/subagent_runtime.py @@ -75,7 +75,9 @@ class ResolvedSubagent: agent_id: UUID | None = None agent_version_id: UUID | None = None collection_names: tuple[str, ...] = () - """The knowledge collections *this* delegate may search. + collection_ids: tuple[UUID, ...] = () + """The knowledge collections *this* delegate may search, and the id of the + knowledge base each name was authorized as, aligned by index. Carried as data beside the agent, rather than left on the deps the build produced, because the library decides what deps a delegation runs with: it diff --git a/backend/app/api/deps.py b/backend/app/api/deps.py index 1f190bb0e..92dcc93d5 100644 --- a/backend/app/api/deps.py +++ b/backend/app/api/deps.py @@ -919,6 +919,8 @@ async def verify_api_key( from app.services.rag.embeddings import EmbeddingService from app.services.rag.ingestion import IngestionService from app.services.rag.documents import DocumentProcessor +from app.services.knowledge_search import KnowledgeSearchService +from app.services.rag.reranker import build_reranker from app.services.rag.retrieval import RetrievalService from app.services.rag.vectorstore import BaseVectorStore, process_vector_store @@ -966,13 +968,31 @@ def get_organization_teardown_service( def get_retrieval_service(vector_store: VectorStoreSvc) -> RetrievalService: - """Create RetrievalService instance.""" - return RetrievalService(vector_store=vector_store, settings=settings.rag) + """Create RetrievalService instance. + + The reranker resolver is `build_reranker`, the one composition point shared + with the agent-run knowledge tool, so both paths rerank the same way. + """ + return RetrievalService( + vector_store=vector_store, + settings=settings.rag, + reranker_resolver=build_reranker, + ) RetrievalSvc = Annotated[RetrievalService, Depends(get_retrieval_service)] +def get_knowledge_search_service( + db: DBSession, retrieval: RetrievalSvc, access: CollectionAccessSvc +) -> KnowledgeSearchService: + """Create KnowledgeSearchService instance.""" + return KnowledgeSearchService(db, retrieval, access) + + +KnowledgeSearchSvc = Annotated[KnowledgeSearchService, Depends(get_knowledge_search_service)] + + def get_document_processor() -> DocumentProcessor: """Create DocumentProcessor instance.""" return DocumentProcessor(settings=settings.rag) @@ -984,9 +1004,14 @@ def get_document_processor() -> DocumentProcessor: def get_ingestion_service( processor: DocumentProcessorSvc, vector_store: VectorStoreSvc, + org: ActiveOrg, ) -> IngestionService: - """Create IngestionService instance.""" - return IngestionService(processor=processor, vector_store=vector_store) + """Create IngestionService instance, scoped to the active organization. + + The organization is what lets the store resolve this tenant's embedding key + on a collection name another tenant may share, rather than another's (#913). + """ + return IngestionService(processor=processor, vector_store=vector_store, organization_id=org.id) IngestionSvc = Annotated[IngestionService, Depends(get_ingestion_service)] diff --git a/backend/app/api/routes/v1/rag.py b/backend/app/api/routes/v1/rag.py index d0b60fa2e..1d1124712 100644 --- a/backend/app/api/routes/v1/rag.py +++ b/backend/app/api/routes/v1/rag.py @@ -48,9 +48,9 @@ CurrentAppAdmin, IngestionSvc, KnowledgeBaseSvc, + KnowledgeSearchSvc, RAGDocumentSvc, RAGSyncSvc, - RetrievalSvc, SyncSourceSvc, VectorStoreSvc, require, @@ -168,7 +168,7 @@ async def create_collection( is refused rather than quietly aliased onto their vector table. """ await access.claim(ctx, name) - await vector_store.create_collection(name) + await vector_store.create_collection(name, organization_id=ctx.organization_id) await kb_svc.create_for_rag_collection( name, user_id=ctx.subject_id, organization_id=ctx.organization_id ) @@ -225,7 +225,9 @@ async def get_collection_info( ) -> Any: """Retrieve stats for a specific collection.""" collection = await access.readable(ctx, name) - return await vector_store.get_collection_info(collection.collection_name) + return await vector_store.get_collection_info( + collection.collection_name, organization_id=collection.organization_id + ) @router.get( @@ -251,33 +253,17 @@ async def list_documents( ) async def search_documents( request: RAGSearchRequest, - retrieval_service: RetrievalSvc, - access: CollectionAccessSvc, + service: KnowledgeSearchSvc, ctx: Auth, ) -> Any: """Search for relevant document chunks. Supports multi-collection search. Every collection named is resolved before the first vector is read, and one the caller cannot reach refuses the whole search rather than being dropped - from it - see `CollectionAccessService.readable_all`. + from it. The embedding and any rerank the search runs are metered against + the caller's organization - see `KnowledgeSearchService`. """ - names = request.collection_names or [request.collection_name] - collections = [kb.collection_name for kb in await access.readable_all(ctx, names)] - if len(collections) > 1: - results = await retrieval_service.retrieve_multi( - query=request.query, - collection_names=collections, - limit=request.limit, - min_score=request.min_score, - ) - else: - results = await retrieval_service.retrieve( - query=request.query, - collection_name=collections[0], - limit=request.limit, - min_score=request.min_score, - filter=request.filter or "", - ) + results = await service.search(ctx, request) api_results = [RAGSearchResult(**hit.model_dump()) for hit in results] return RAGSearchResponse(results=api_results) diff --git a/backend/app/commands/rag.py b/backend/app/commands/rag.py index b5dbf3a7a..f56e958b3 100644 --- a/backend/app/commands/rag.py +++ b/backend/app/commands/rag.py @@ -72,7 +72,11 @@ def get_rag_services() -> tuple[ vector_store = process_vector_store(settings, embedder) processor = DocumentProcessor(settings=settings) retrieval = RetrievalService(vector_store=vector_store, settings=settings) - ingestion = IngestionService(processor=processor, vector_store=vector_store) + # CLI admin context: no single tenant, so resolution falls back to the name. + # Uploads through the API are org-scoped. + ingestion = IngestionService( + processor=processor, vector_store=vector_store, organization_id=None + ) return settings, vector_store, processor, retrieval, ingestion @@ -92,7 +96,7 @@ async def list_collections_async(vector_store: BaseVectorStore) -> None: for name in collection_names: try: - info_obj = await vector_store.get_collection_info(name) + info_obj = await vector_store.get_collection_info(name, organization_id=None) click.echo(f" {name}") click.echo(f" Vectors: {info_obj.total_vectors:,}") click.echo(f" Dimension: {info_obj.dim}") @@ -343,6 +347,7 @@ async def search_async( query=query, collection_name=collection, limit=top_k, + organization_id=None, ) if not results: @@ -473,7 +478,7 @@ async def stats_async(settings: RAGSettings, vector_store: BaseVectorStore) -> N total_vectors = 0 for name in collection_names: try: - info_obj = await vector_store.get_collection_info(name) + info_obj = await vector_store.get_collection_info(name, organization_id=None) click.echo(f" {name}:") click.echo(f" Vectors: {info_obj.total_vectors:,}") total_vectors += info_obj.total_vectors diff --git a/backend/app/db/models/ingestion_spend.py b/backend/app/db/models/ingestion_spend.py index 3f185be26..21239ebf2 100644 --- a/backend/app/db/models/ingestion_spend.py +++ b/backend/app/db/models/ingestion_spend.py @@ -1,21 +1,25 @@ -"""What indexing a document cost - embedding spend outside any agent run. +"""RAG spend that lands on no agent run - indexing, and a direct search. Agent runs record their cost on `agent_runs.cost_usd`, and that includes the -embeddings a knowledge search makes, because a run's ledger is metering while -it executes. Ingestion has no run: a document is embedded in a worker, on -nobody's conversation, and for months that spend was recorded nowhere - an -organization could embed unbounded volume under an exhausted budget, because -the monthly total only ever summed runs. - -One row per model per metering window - a document upload, a connector sync - -rather than per API call. The unit someone reconciles against a bill is "what -did indexing this cost with which model", not "what did chunk 37 cost"; and a -window can spend in two models at once, because describing a scanned page is a -vision call and embedding it is not. +embeddings a knowledge search *inside a run* makes, because a run's ledger is +metering while it executes. Two RAG activities have no run to bill: indexing a +document happens in a worker on nobody's conversation, and `POST /rag/search` +answers a caller directly. For months the first was recorded nowhere - an +organization could embed unbounded volume under an exhausted budget, because the +monthly total only ever summed runs - and the second, once metered, landed here +too. `source` tells them apart (:class:`SpendSource`) so the dashboard does not +report a search as indexing; both still count toward the monthly budget. + +One row per model per metering window - a document upload, a connector sync, one +search - rather than per API call. The unit someone reconciles against a bill is +"what did this cost with which model", not "what did chunk 37 cost"; and a window +can spend in two models at once, because describing a scanned page is a vision +call and embedding it is not. """ import uuid from decimal import Decimal +from enum import StrEnum from sqlalchemy import ForeignKey, Index, Integer, Numeric, String from sqlalchemy.dialects.postgresql import UUID as PG_UUID @@ -24,6 +28,20 @@ from app.db.base import Base, TimestampMixin +class SpendSource(StrEnum): + """Which RAG activity a row of non-run spend paid for. + + The table began as indexing alone, then a metered knowledge search landed + its embedding and rerank cost here too - both are RAG spend outside any + agent run. Left undistinguished, a search inflated the dashboard's + "indexing" subtotal, so this says which is which. Both still count toward + the monthly budget; only the reporting split reads the column. + """ + + INGESTION = "ingestion" + RETRIEVAL = "retrieval" + + class IngestionSpend(Base, TimestampMixin): __tablename__ = "ingestion_spend" @@ -59,6 +77,15 @@ class IngestionSpend(Base, TimestampMixin): # True when the model had no price - the cost is then a floor. cost_is_partial: Mapped[bool] = mapped_column(nullable=False, default=False) + # Indexing or retrieval. `server_default` so the rows written before the + # column existed - all of them indexing - read as such without a backfill. + source: Mapped[str] = mapped_column( + String(16), + nullable=False, + default=SpendSource.INGESTION.value, + server_default=SpendSource.INGESTION.value, + ) + # Declared here as well as in the migration: the integration tests build # the schema from the models, and the monthly lookup queries exactly this # shape - one organization, one window. diff --git a/backend/app/db/models/knowledge_base.py b/backend/app/db/models/knowledge_base.py index 69d56b72c..c871461e4 100644 --- a/backend/app/db/models/knowledge_base.py +++ b/backend/app/db/models/knowledge_base.py @@ -70,6 +70,19 @@ class KnowledgeBase(TimestampMixin, Base): ForeignKey("organization_secrets.id", ondelete="SET NULL"), nullable=True, ) + # Reranking is on for this collection only when both are set. The model is + # the reranker's name (e.g. 'rerank-v3.5'); the secret is the org vault key + # that pays for it. NULL on either means no reranking - retrieval behaves + # exactly as it did before the feature. SET NULL on delete for the same + # reason as the embedding key: losing the key drops reranking, it does not + # take search down. There is no deployment fallback, because a reranker with + # no key is simply off. + rerank_model: Mapped[str | None] = mapped_column(String(128), nullable=True) + rerank_secret_id: Mapped[uuid.UUID | None] = mapped_column( + PG_UUID(as_uuid=True), + ForeignKey("organization_secrets.id", ondelete="SET NULL"), + nullable=True, + ) # How widely the collection is exposed inside its org; combines with the # member's role scope and any explicit grant (app.services.access). visibility: Mapped[str] = mapped_column( diff --git a/backend/app/repositories/ingestion_spend.py b/backend/app/repositories/ingestion_spend.py index 7790191a2..b92c497de 100644 --- a/backend/app/repositories/ingestion_spend.py +++ b/backend/app/repositories/ingestion_spend.py @@ -7,7 +7,7 @@ from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession -from app.db.models.ingestion_spend import IngestionSpend +from app.db.models.ingestion_spend import IngestionSpend, SpendSource async def record( @@ -20,6 +20,7 @@ async def record( output_tokens: int, cost_usd: Decimal, cost_is_partial: bool, + source: SpendSource = SpendSource.INGESTION, ) -> IngestionSpend: spend = IngestionSpend( organization_id=organization_id, @@ -29,6 +30,7 @@ async def record( output_tokens=output_tokens, cost_usd=cost_usd, cost_is_partial=cost_is_partial, + source=source.value, ) db.add(spend) await db.flush() @@ -37,7 +39,11 @@ async def record( async def sum_cost_since(db: AsyncSession, *, organization_id: UUID, since: datetime) -> Decimal: - """Total ingestion spend in a window - the half of a monthly budget runs cannot see.""" + """Total non-run RAG spend in a window - the half of a monthly budget runs cannot see. + + Every source: indexing and a direct search both count toward the cap, so + this is deliberately unfiltered where the dashboard split is not. + """ result = await db.scalar( select(func.coalesce(func.sum(IngestionSpend.cost_usd), 0)).where( IngestionSpend.organization_id == organization_id, @@ -48,21 +54,32 @@ async def sum_cost_since(db: AsyncSession, *, organization_id: UUID, since: date async def sum_cost_window( - db: AsyncSession, *, organization_id: UUID, start: datetime, end: datetime + db: AsyncSession, + *, + organization_id: UUID, + start: datetime, + end: datetime, + source: SpendSource | None = None, ) -> Decimal: - """Ingestion spend inside a closed window - the dashboard's half of the bill. + """Non-run RAG spend inside a closed window - the dashboard's half of the bill. Distinct from :func:`sum_cost_since`, which is open-ended and feeds budget enforcement: a cap is measured against the calendar month, a dashboard period against whatever window its filter chose. Half-open at the end, the same way `agent_run_repo.sum_cost_window` is, so a document indexed at 23:59:59 on the last day counts once and only once. + + `source` narrows to indexing or retrieval; the dashboard sums each on its + own so a search is not reported as indexing. `None` sums both. """ + conditions = [ + IngestionSpend.organization_id == organization_id, + IngestionSpend.created_at >= start, + IngestionSpend.created_at < end, + ] + if source is not None: + conditions.append(IngestionSpend.source == source.value) result = await db.scalar( - select(func.coalesce(func.sum(IngestionSpend.cost_usd), 0)).where( - IngestionSpend.organization_id == organization_id, - IngestionSpend.created_at >= start, - IngestionSpend.created_at < end, - ) + select(func.coalesce(func.sum(IngestionSpend.cost_usd), 0)).where(*conditions) ) return Decimal(result or 0) diff --git a/backend/app/repositories/knowledge_base.py b/backend/app/repositories/knowledge_base.py index 13a03bf05..339d137fa 100644 --- a/backend/app/repositories/knowledge_base.py +++ b/backend/app/repositories/knowledge_base.py @@ -94,6 +94,8 @@ async def create( organization_id: UUID | None = None, is_default: bool = False, embedding_secret_id: UUID | None = None, + rerank_model: str | None = None, + rerank_secret_id: UUID | None = None, visibility: str | None = None, ) -> KnowledgeBase: """Create a knowledge base. @@ -119,6 +121,8 @@ async def create( embedding_dim=embedding_dim, embedding_provider=embedding_provider, embedding_secret_id=embedding_secret_id, + rerank_model=rerank_model, + rerank_secret_id=rerank_secret_id, **({"visibility": visibility} if visibility is not None else {}), ) db.add(kb) @@ -134,6 +138,9 @@ async def update( name: str | None = None, description: str | None = None, ingestion_config: dict[str, object] | None = None, + set_rerank: bool = False, + rerank_model: str | None = None, + rerank_secret_id: UUID | None = None, embedding_provider: str | None = None, embedding_secret_id: UUID | None = None, clear_embedding_secret: bool = False, @@ -150,6 +157,12 @@ async def update( db_kb.description = description if ingestion_config is not None: db_kb.ingestion_config = ingestion_config + # A pair set together, and the only field here that can be set back to null: + # `set_rerank` is what tells "turn reranking off" from "leave it alone", + # which the None-means-skip convention above cannot express. + if set_rerank: + db_kb.rerank_model = rerank_model + db_kb.rerank_secret_id = rerank_secret_id if embedding_provider is not None: db_kb.embedding_provider = embedding_provider if clear_embedding_secret: @@ -210,3 +223,50 @@ async def list_by_collection_name(db: AsyncSession, collection_name: str) -> lis .order_by(KnowledgeBase.created_at) ) return list(result.scalars().all()) + + +async def knowledge_bases_using( + db: AsyncSession, *, organization_id: UUID, secret_id: UUID +) -> list[tuple[UUID, str]]: + """Knowledge bases that reference this secret as their embedding or rerank key. + + Both columns in one query: a key is bound through `embedding_secret_id` or + `rerank_secret_id`, and either binding breaks the same way when the key is + deleted - the foreign key nulls the reference (SET NULL) and the collection + silently stops embedding or reranking. A vault listing that only checks + agent specs calls such a key unused and invites exactly that deletion, so + this is what lets the listing account for the collections too. Scoped to the + organization, like every other lookup here. + """ + result = await db.execute( + select(KnowledgeBase.id, KnowledgeBase.name) + .where( + KnowledgeBase.organization_id == organization_id, + or_( + KnowledgeBase.embedding_secret_id == secret_id, + KnowledgeBase.rerank_secret_id == secret_id, + ), + ) + .order_by(KnowledgeBase.name) + ) + return [(row[0], row[1]) for row in result.all()] + + +async def get_for_collection( + db: AsyncSession, collection_name: str, organization_id: UUID | None +) -> KnowledgeBase | None: + """The knowledge base an organization resolves a collection name to. + + `collection_name` is not unique across tenants, so resolving one by name + alone can return another organization's row - and then unseal and bill that + organization's key (#913). The organization narrows the candidates in two + passes: its own row wins, and an `app`-scoped collection (owned by no + organization) is the shared fallback. `organization_id` is `None` only where + there is genuinely no tenant to scope to - a CLI ingest - and then the first + candidate stands, which is the old name-only behaviour for that path alone. + """ + candidates = await list_by_collection_name(db, collection_name) + for kb in candidates: + if organization_id is None or kb.organization_id == organization_id: + return kb + return next((kb for kb in candidates if kb.organization_id is None), None) diff --git a/backend/app/schemas/knowledge_base.py b/backend/app/schemas/knowledge_base.py index e4b820920..03bc0ed48 100644 --- a/backend/app/schemas/knowledge_base.py +++ b/backend/app/schemas/knowledge_base.py @@ -50,6 +50,22 @@ class KnowledgeBaseCreate(BaseSchema): "can be paid with." ), ) + rerank_model: str | None = Field( + default=None, + max_length=128, + description=( + "Which reranker reorders this collection's search results. Reranking " + "runs only when this and rerank_secret_id are both set; omit both to " + "leave it off. Unlike the embedding model this can be changed later." + ), + ) + rerank_secret_id: UUID | None = Field( + default=None, + description=( + "The organization vault key (a Cohere key) that pays for reranking. " + "Set together with rerank_model." + ), + ) ingestion_config: IngestionConfig | None = Field( default=None, description=( @@ -81,6 +97,11 @@ class KnowledgeBaseUpdate(BaseSchema): "documents ingested afterwards; nothing already indexed is re-parsed." ), ) + # Sent as a pair: both to turn reranking on or change it, both null to turn + # it off. Whether the caller meant to touch reranking at all is read from + # the fields they actually sent, so an update that omits both leaves it be. + rerank_model: str | None = Field(default=None, max_length=128) + rerank_secret_id: UUID | None = Field(default=None) embedding_provider: str | None = Field( default=None, max_length=32, @@ -128,6 +149,10 @@ class KnowledgeBaseRead(BaseSchema, TimestampSchema): # Editable, unlike the two above - see `KnowledgeBaseUpdate`. embedding_provider: str embedding_secret_id: UUID | None = None + # Both null unless reranking is configured; the secret id is safe to expose, + # it names a vault row rather than carrying its value. + rerank_model: str | None = None + rerank_secret_id: UUID | None = None # Derived per request from `rag_documents`, not stored. Defaulted rather than # required so the single-row responses - create, read, update - stay # constructible straight from the ORM row, which is what they are: a diff --git a/backend/app/schemas/secret.py b/backend/app/schemas/secret.py index c030a94e6..eab1116ac 100644 --- a/backend/app/schemas/secret.py +++ b/backend/app/schemas/secret.py @@ -80,7 +80,7 @@ class SecretPurposeList(BaseSchema): class SecretUsage(BaseSchema): """One place a stored key is bound. Named so the answer is readable.""" - kind: Literal["agent"] + kind: Literal["agent", "knowledge_base"] id: UUID name: str diff --git a/backend/app/schemas/stats.py b/backend/app/schemas/stats.py index 0d3c1ee8c..bf2417303 100644 --- a/backend/app/schemas/stats.py +++ b/backend/app/schemas/stats.py @@ -111,15 +111,19 @@ class ProviderCost(BaseSchema): class CostBlock(BaseSchema): """The window's whole bill, and the two halves it is made of. - `period_usd` is models **plus ingestion**, which is the same arithmetic - `spend.organization_spend_since` measures a monthly cap on. It used to be - models alone, so the dashboard's headline and the month-to-date line under - it were two different definitions of cost sitting on one card with nothing - saying so - on a deployment that indexes documents they simply disagreed. - - `model_usd` and `ingestion_usd` sum to it. They are separate fields rather - than a computed split because the two are answered by different tables and - a reader deciding where the money went should not have to subtract. + `period_usd` is models **plus ingestion plus retrieval**, which is the same + arithmetic `spend.organization_spend_since` measures a monthly cap on. It + used to be models alone, so the dashboard's headline and the month-to-date + line under it were two different definitions of cost sitting on one card with + nothing saying so - on a deployment that indexes documents they simply + disagreed. + + `model_usd`, `ingestion_usd` and `retrieval_usd` sum to it. They are separate + fields rather than a computed split because they are answered by different + tables and sources, and a reader deciding where the money went should not + have to subtract. `retrieval_usd` is what a metered `POST /rag/search` spent + on embeddings and reranking - kept apart from `ingestion_usd` so a search is + not reported as indexing. `previous_period_usd` is the whole bill too, so the change against the last window compares like with like. @@ -132,6 +136,7 @@ class CostBlock(BaseSchema): previous_period_usd: Decimal model_usd: Decimal ingestion_usd: Decimal + retrieval_usd: Decimal by_provider: list[ProviderCost] diff --git a/backend/app/services/agent_runner.py b/backend/app/services/agent_runner.py index 3f0cdb201..e57aac913 100644 --- a/backend/app/services/agent_runner.py +++ b/backend/app/services/agent_runner.py @@ -1318,7 +1318,12 @@ def build(*, name: str, instructions: str, model: str) -> PydanticAgent[Any, Any ), model=profiles[model], agent_id=delegation.agent_id, - resources={"kb_collection_names": [], "skills": [], CONTEXT_FILES_RESOURCE: []}, + resources={ + "kb_collection_names": [], + "kb_collection_ids": [], + "skills": [], + CONTEXT_FILES_RESOURCE: [], + }, secrets={}, extra_toolsets=[], )() @@ -1424,13 +1429,16 @@ def __init__(self, db: AsyncSession) -> None: self.proposals = SkillProposalService(db) self.transcript = TranscriptService(db) - async def _collection_names(self, spec: AgentSpec, ctx: AuthContext) -> list[str]: - """Vector-store collection names for the agent's bound collections. + async def _bound_collections(self, spec: AgentSpec, ctx: AuthContext) -> list[tuple[str, UUID]]: + """The agent's bound collections as (collection_name, knowledge_base_id). Resolved server-side and passed through deps: the model asks *what* to - search, never *where*. + search, never *where*. The id travels beside the name because + `collection_name` is not unique - search resolves each collection by the + bound id rather than re-selecting by name, so a name shared with another + knowledge base cannot resolve that row's config or key (#913). """ - names: list[str] = [] + bound: list[tuple[str, UUID]] = [] for collection_id in spec.collection_ids: collection = await knowledge_base_repo.get_by_id(self.db, collection_id) if collection is None or collection.organization_id != ctx.organization_id: @@ -1442,8 +1450,8 @@ async def _collection_names(self, spec: AgentSpec, ctx: AuthContext) -> list[str ctx.organization_id, ) continue - names.append(collection.collection_name) - return names + bound.append((collection.collection_name, collection.id)) + return bound async def _recorded_conversation_state(self, conversation_id: UUID | None) -> _RecordedState: """What an earlier turn of this thread recorded, in one read. @@ -1623,8 +1631,10 @@ async def _assemble( # Everything a capability needs but must not fetch itself. Resolved once, # server-side, so the model cannot influence what an agent reaches. + bound_collections = await self._bound_collections(spec, ctx) resources: dict[str, Any] = { - "kb_collection_names": await self._collection_names(spec, ctx), + "kb_collection_names": [name for name, _ in bound_collections], + "kb_collection_ids": [cid for _, cid in bound_collections], "skills": await self.skills.resolve_for_agent(ctx, spec.skill_ids), CONTEXT_FILES_RESOURCE: await self.context.resolve_for_agent(ctx, spec.context_ids), } @@ -2139,6 +2149,7 @@ async def _resolve_specialist( max_steps=specialist.max_steps, preferred_mode=specialist.preferred_mode, collection_names=tuple(own_resources["kb_collection_names"]), + collection_ids=tuple(own_resources["kb_collection_ids"]), ) async def _resolve_delegate( @@ -2318,6 +2329,7 @@ async def _resolve_delegate( # replaces those deps with a clone of the parent's - see # `ResolvedSubagent.collection_names`. collection_names=tuple(delegate_resources["kb_collection_names"]), + collection_ids=tuple(delegate_resources["kb_collection_ids"]), ) async def _delegate_resources( @@ -2353,8 +2365,10 @@ async def _delegate_resources( workspace is opened per delegate: only the run has one. Sharing is how a delegate reaches a durable workspace at all. """ + bound_collections = await self._bound_collections(spec, ctx) resources: dict[str, Any] = { - "kb_collection_names": await self._collection_names(spec, ctx), + "kb_collection_names": [name for name, _ in bound_collections], + "kb_collection_ids": [cid for _, cid in bound_collections], "skills": await self.skills.resolve_for_agent(ctx, spec.skill_ids), CONTEXT_FILES_RESOURCE: await self.context.resolve_for_agent(ctx, spec.context_ids), } diff --git a/backend/app/services/embedding_resolution.py b/backend/app/services/embedding_resolution.py index 795a8cc0c..f5d4c2e4a 100644 --- a/backend/app/services/embedding_resolution.py +++ b/backend/app/services/embedding_resolution.py @@ -35,6 +35,7 @@ import logging from dataclasses import dataclass from enum import StrEnum +from uuid import UUID from sqlalchemy.ext.asyncio import AsyncSession @@ -151,7 +152,9 @@ def describe(self, collection_name: str) -> str: ) -async def embeddings_for_collection(collection_name: str) -> ResolvedEmbeddings | None: +async def embeddings_for_collection( + collection_name: str, organization_id: UUID | None, knowledge_base_id: UUID | None = None +) -> ResolvedEmbeddings | None: """Resolve one collection's embedding model, provider and credential. Returns None for a collection no knowledge base claims - the store then @@ -159,6 +162,15 @@ async def embeddings_for_collection(collection_name: str) -> ResolvedEmbeddings gotten. Opens its own session because the store embeds from places with no request in sight: a worker mid-ingestion, a capability mid-run. + `knowledge_base_id`, when given, is the knowledge base the caller was already + authorized against, and resolution reads *that* row. `collection_name` is not + unique, so a name+organization lookup can return a different row than the + access check authorized - a restricted `org` collection sharing an `app` + collection's name - and then unseal and bill a key the caller was never + granted (#913). The search path passes the authorized id; ingestion and the + CLI, which choose the row themselves, pass none and fall back to the + `organization_id`-scoped lookup. + A provider the catalog no longer names - an entry removed from the file under a collection that was using it - resolves to the deployment's, with a log line. The alternative is a collection nobody can search because a @@ -166,7 +178,11 @@ async def embeddings_for_collection(collection_name: str) -> ResolvedEmbeddings address this build is certain of. """ async with get_db_context() as db: - kb = await knowledge_base_repo.get_by_collection_name(db, collection_name) + kb = ( + await knowledge_base_repo.get_by_id(db, knowledge_base_id) + if knowledge_base_id is not None + else await knowledge_base_repo.get_for_collection(db, collection_name, organization_id) + ) if kb is None: return None provider = embedding_providers.get(kb.embedding_provider) diff --git a/backend/app/services/knowledge_base.py b/backend/app/services/knowledge_base.py index 1fc08da5b..bff559ce4 100644 --- a/backend/app/services/knowledge_base.py +++ b/backend/app/services/knowledge_base.py @@ -34,6 +34,7 @@ deployment_embedding, ) from app.services.rag import embedding_providers +from app.services.rerank_resolution import RERANK_KEY_PURPOSES, SUPPORTED_RERANK_MODELS logger = logging.getLogger(__name__) @@ -314,6 +315,9 @@ async def create( await self._check_embedding_secret( data.embedding_secret_id, ctx=ctx, organization_id=org_id, provider=provider ) + self._check_rerank_pair(data.rerank_model, data.rerank_secret_id) + if data.rerank_secret_id is not None: + await self._check_rerank_secret(ctx, data.rerank_secret_id, organization_id=org_id) return await knowledge_base_repo.create( self.db, name=data.name, @@ -327,6 +331,8 @@ async def create( embedding_dim=embedding_dim, embedding_provider=provider.provider, embedding_secret_id=data.embedding_secret_id, + rerank_model=data.rerank_model, + rerank_secret_id=data.rerank_secret_id, ) async def _check_embedding_secret( @@ -374,6 +380,74 @@ async def _check_embedding_secret( purpose=row.purpose, ) + @staticmethod + def _check_rerank_pair(model: str | None, secret_id: UUID | None) -> None: + """A reranker is a supported model *and* a key, or neither. + + Reranking runs only when both are set (`rerank_resolution`), so a lone + half is a setting that reads as configured and does nothing. And a model + this deployment cannot run is the same failure by another route: it is + accepted, stored and shown as configured, then every search fails inside + Cohere and is swallowed, so reranking is silently off. Both are refused + here, where the person setting it can see why, rather than at search time. + """ + if (model is None) != (secret_id is None): + raise BadRequestError( + message="Reranking needs both a model and a key, or neither", + details={"rerank_model": model, "rerank_secret_id": str(secret_id)}, + ) + if model is not None and model not in SUPPORTED_RERANK_MODELS: + raise BadRequestError( + message=( + f"Unsupported rerank model; this deployment reranks through " + f"{', '.join(SUPPORTED_RERANK_MODELS)}" + ), + details={"rerank_model": model}, + ) + + async def _check_rerank_secret( + self, ctx: AuthContext, secret_id: UUID, *, organization_id: UUID | None + ) -> None: + """Refuse a rerank key the caller may not use, or one of the wrong kind. + + Checked at creation and on update, where the person choosing can fix it - + resolution degrades a bad key to no reranking, so this is the one moment + a wrong choice is visible. + + Binding a key is lending it: reranking spends it for everyone who can + search the collection, so whoever sets it has to be able to reach the key + themselves. The picker only ever offers what they can see, but the API + took an id, and a private key another member owns is in the organization's + vault yet not theirs to use - so an org-scoped lookup alone would let a + `collections:edit` holder bind a key `secrets:view` would refuse them. + Refused as "not in the vault", the same as a genuine miss, so a refusal + cannot be told apart and used to enumerate it - exactly as agent secret + bindings are checked (`agent_registry`). + """ + if organization_id is None: + raise BadRequestError( + message="Only an organization collection can carry a vault key", + details={"rerank_secret_id": str(secret_id)}, + ) + row = await organization_secret_repo.get( + self.db, secret_id, organization_id=organization_id + ) + if row is None or not await resolve_access( + self.db, ctx, row, Perm.SECRETS_VIEW, resource_type=SECRET + ): + raise BadRequestError( + message="That key is not in this organization's vault", + details={"rerank_secret_id": str(secret_id)}, + ) + if row.purpose not in RERANK_KEY_PURPOSES: + raise BadRequestError( + message=( + f"That key is for {row.purpose}; reranking runs through " + f"{', '.join(RERANK_KEY_PURPOSES)}" + ), + details={"purpose": row.purpose}, + ) + async def update( self, kb_id: UUID, @@ -387,6 +461,17 @@ async def update( if data.ingestion_config is None else await self._usable_config(ctx, data.ingestion_config) ) + # The rerank pair is set only when the caller actually sent it, so an + # update about something else leaves reranking untouched; sending both + # as null is how it is turned off, which is why the trigger is "was the + # field present" rather than "is it not None". + sets_rerank = bool({"rerank_model", "rerank_secret_id"} & data.model_fields_set) + if sets_rerank: + self._check_rerank_pair(data.rerank_model, data.rerank_secret_id) + if data.rerank_secret_id is not None: + await self._check_rerank_secret( + ctx, data.rerank_secret_id, organization_id=kb.organization_id + ) # The provider the collection will be on when this update lands, which is # what the key has to match: moving to OpenAI and choosing an OpenAI key # in one request must be accepted, and either half alone must be checked @@ -413,6 +498,9 @@ async def update( name=data.name, description=data.description, ingestion_config=None if config is None else config.model_dump(mode="json"), + set_rerank=sets_rerank, + rerank_model=data.rerank_model, + rerank_secret_id=data.rerank_secret_id, embedding_provider=data.embedding_provider, embedding_secret_id=data.embedding_secret_id, clear_embedding_secret=data.clear_embedding_secret, diff --git a/backend/app/services/knowledge_search.py b/backend/app/services/knowledge_search.py new file mode 100644 index 000000000..a83a5ef3e --- /dev/null +++ b/backend/app/services/knowledge_search.py @@ -0,0 +1,151 @@ +"""The knowledge-search request path, and the spend it books. + +`POST /rag/search` used to be the one metered gap in RAG: it embedded the query +and, once configured, reranks - both of which cost money - inside no +`metered_by` block and against no ledger, so neither landed on the +organization's monthly bill (the #16 class of defect). This service closes that. + +It opens a ledger scoped to the caller's organization, runs the search inside a +`metered_by` block so the ambient embedding and rerank calls book to it, and +persists what they spent to `ingestion_spend` - the same sink a worker's +ingestion spend lands in, with a null document id because a search indexed +nothing. Reranking is what made the gap worth closing; metering the embeddings +too is the beneficial side effect. + +The search itself is unchanged: the same collection-access resolution and the +same single- vs multi-collection retrieval the route did inline, moved behind a +service so the route stays HTTP plumbing. +""" + +from __future__ import annotations + +from decimal import Decimal +from typing import TYPE_CHECKING + +from app.agents.capabilities.budget import SpendLedger, metered_by +from app.db.models.ingestion_spend import SpendSource +from app.db.session import get_db_context +from app.repositories import ingestion_spend_repo +from app.services.spend import assert_organization_within_budget + +if TYPE_CHECKING: + from sqlalchemy.ext.asyncio import AsyncSession + + from app.core.permissions import AuthContext + from app.schemas.rag import RAGSearchRequest + from app.services.collection_access import CollectionAccessService + from app.services.rag.models import SearchResult + from app.services.rag.retrieval import RetrievalService + + +class KnowledgeSearchService: + """Resolves, meters and runs a knowledge search on the request path.""" + + def __init__( + self, + db: AsyncSession, + retrieval: RetrievalService, + access: CollectionAccessService, + ) -> None: + self.db = db + self.retrieval = retrieval + self.access = access + + async def search(self, ctx: AuthContext, request: RAGSearchRequest) -> list[SearchResult]: + """Run the search, metered against the caller's organization. + + A collection the caller cannot reach refuses the whole search rather + than being dropped from it - `CollectionAccessService.readable_all` + raises, and this never opens a ledger for a search that will not run. + + The budget is asserted before the ledger opens, not only recorded after: + the embedding and any rerank spend this path incurs is real provider + cost, and without the same guard ingestion carries a member with + collection-view access could keep spending the organization's keys after + the monthly cap is reached. Checked here, once access is resolved, so a + refused search never reaches a paid call. + """ + names = request.collection_names or [request.collection_name] + # Keep the authorized rows, not just their names: `collection_name` is + # not unique, so resolution must read the exact knowledge base access + # granted - passing the name back would let it re-select a different + # same-named row and unseal that row's key (#913). + authorized = await self.access.readable_all(ctx, names) + collections = [kb.collection_name for kb in authorized] + knowledge_base_ids = [kb.id for kb in authorized] + + if ctx.organization_id is not None: + await assert_organization_within_budget(self.db, ctx.organization_id) + + ledger = SpendLedger(organization_id=ctx.organization_id) + try: + with metered_by(ledger): + if len(collections) > 1: + results = await self.retrieval.retrieve_multi( + query=request.query, + collection_names=collections, + limit=request.limit, + min_score=request.min_score, + organization_id=ctx.organization_id, + knowledge_base_ids=knowledge_base_ids, + ) + else: + results = await self.retrieval.retrieve( + query=request.query, + collection_name=collections[0], + limit=request.limit, + min_score=request.min_score, + filter=request.filter or "", + organization_id=ctx.organization_id, + knowledge_base_id=knowledge_base_ids[0], + ) + except Exception: + # The query embedding is booked before the vector query it pays for, + # so a search that fails mid-flight has already spent. Recording that + # on the request session would be undone with the failed request's + # rollback, so the failure path books through a session of its own + # that commits: the platform records spend even when the run fails. + await self._book_failed_spend(ledger) + raise + + await self._record_spend(self.db, ledger) + return results + + async def _book_failed_spend(self, ledger: SpendLedger) -> None: + """Persist a failed search's spend in a transaction of its own. + + The request that raised is about to roll back, taking `self.db` with it, + so what was already spent is written through a fresh session that commits + independently. Skips opening one when nothing was spent. + """ + if not ledger.entries: + return + async with get_db_context() as db: + await self._record_spend(db, ledger) + + @staticmethod + async def _record_spend(db: AsyncSession, ledger: SpendLedger) -> None: + """Persist what the search spent, one row per model, if it spent anything. + + A null `rag_document_id` because a search indexes no document; the + organization is what a monthly budget reads this back against. Priced by + the reranker and by embeddings independently, so a partial cost is one + the embedding half could not price, exactly as ingestion records it. + Tagged `RETRIEVAL` so the dashboard reports it as search, not indexing; + it still counts toward the monthly budget alongside indexing. + """ + if not ledger.entries: + return + for model in dict.fromkeys(entry.model_name for entry in ledger.entries): + entries = [entry for entry in ledger.entries if entry.model_name == model] + await ingestion_spend_repo.record( + db, + organization_id=ledger.organization_id, + rag_document_id=None, + model=model, + input_tokens=sum(entry.input_tokens for entry in entries), + output_tokens=sum(entry.output_tokens for entry in entries), + cost_usd=sum((entry.cost_usd for entry in entries), Decimal(0)), + cost_is_partial=any(not entry.priced for entry in entries), + source=SpendSource.RETRIEVAL, + ) diff --git a/backend/app/services/organization_secret.py b/backend/app/services/organization_secret.py index b320be48d..026a22dbd 100644 --- a/backend/app/services/organization_secret.py +++ b/backend/app/services/organization_secret.py @@ -35,7 +35,12 @@ from app.core.vault import VaultScope from app.db.models.organization_secret import OrganizationSecret from app.db.models.resource_grant import Visibility -from app.repositories import member_repo, organization_secret_repo, resource_grant_repo +from app.repositories import ( + knowledge_base_repo, + member_repo, + organization_secret_repo, + resource_grant_repo, +) from app.schemas.resource_grant import as_visibility from app.schemas.secret import SecretRead, SecretUsage from app.services.access import SECRET, resolve_access, visible_resource_ids @@ -136,6 +141,9 @@ async def list_secrets( used = await organization_secret_repo.agents_using( self.db, organization_id=ctx.organization_id, secret_id=secret.id ) + kbs = await knowledge_base_repo.knowledge_bases_using( + self.db, organization_id=ctx.organization_id, secret_id=secret.id + ) rows.append( SecretRead( id=secret.id, @@ -153,6 +161,10 @@ async def list_secrets( shared_with=shared_counts.get(secret.id, 0), used_by=[ SecretUsage(kind="agent", id=agent_id, name=name) for agent_id, name in used + ] + + [ + SecretUsage(kind="knowledge_base", id=kb_id, name=name) + for kb_id, name in kbs ], created_at=secret.created_at, updated_at=secret.updated_at, diff --git a/backend/app/services/rag/ingestion.py b/backend/app/services/rag/ingestion.py index c4dd278d2..1e6ec5947 100644 --- a/backend/app/services/rag/ingestion.py +++ b/backend/app/services/rag/ingestion.py @@ -4,6 +4,7 @@ from collections.abc import Awaitable, Callable from dataclasses import dataclass from pathlib import Path +from uuid import UUID from app.services.rag.documents import DocumentProcessor from app.services.rag.failures import IngestionStage, failure_summary @@ -44,10 +45,15 @@ def __init__( processor: DocumentProcessor, vector_store: BaseVectorStore, on_event: Callable[..., Awaitable[None]] | None = None, + *, + organization_id: UUID | None, ): self.processor = processor self.store = vector_store self._on_event = on_event + # The organization this ingest embeds for, so the store resolves this + # tenant's key and not another's on a shared collection name (#913). + self._organization_id = organization_id async def _emit(self, event: str, data: dict[str, object]) -> None: if self._on_event: @@ -128,6 +134,7 @@ async def ingest_file( await self.store.insert_document( collection_name=collection_name, document=document, + organization_id=self._organization_id, ) if existing_id: diff --git a/backend/app/services/rag/reranker.py b/backend/app/services/rag/reranker.py new file mode 100644 index 000000000..59afb66ad --- /dev/null +++ b/backend/app/services/rag/reranker.py @@ -0,0 +1,201 @@ +"""Reordering retrieved candidates by a model's judgement, not by distance. + +Vector search sorts by embedding distance, which is a proxy for relevance and +sometimes a poor one. A reranker is a second pass: a cross-encoder scores each +candidate against the query directly and reorders them. Retrieval overfetches, +hands the candidates here, and truncates to the caller's limit afterwards, so a +better answer sitting tenth by distance can surface in the top few. + +Cohere is the only provider today. A second one is a second :class:`BaseReranker` +- the interface is what retrieval depends on, and a collection's resolved +credential is what decides which implementation, if any, it gets +(:mod:`app.services.rerank_resolution`). + +Spend is booked here, not through :func:`record_ambient_usage`: a rerank call is +priced per search, not per token, and `genai-prices` - which prices everything +else this platform meters - does not know rerank models and would book it +`cost_usd=0, priced=False`. So the cost is computed from Cohere's published +per-search price and handed to :func:`book_ambient_spend` already finished, +landing on whichever ledger is metering the search. +""" + +from __future__ import annotations + +import logging +from abc import ABC, abstractmethod +from decimal import Decimal +from math import ceil +from typing import TYPE_CHECKING +from uuid import UUID + +import cohere + +from app.agents.capabilities.budget import SpendEntry, book_ambient_spend +from app.services.rag.models import SearchResult +from app.services.rerank_resolution import reranker_for_collection + +if TYPE_CHECKING: + from cohere import AsyncClientV2 + +logger = logging.getLogger(__name__) + +# Cohere bills reranking per "search unit": one query with up to this many +# documents. The billed figure comes from the response; this is only the +# fallback estimate when the response omits it - one unit per this many +# candidates, which ignores the document-splitting the real figure accounts for. +_DOCS_PER_SEARCH_UNIT = 100 + +# USD per search unit. Cohere Rerank 3.5 is $2.00 per 1,000 searches. +# Confirmed against https://cohere.com/pricing on 2026-08-18; genai-prices does +# not carry rerank models, so this is the one price in the metering path that +# lives in this repository. A maintainer changing the model or seeing the bill +# drift should re-check that page and this number together. +_PRICE_PER_SEARCH_UNIT_USD = Decimal("0.002") + + +class BaseReranker(ABC): + """Reorders retrieved candidates against the query. What retrieval depends on.""" + + @abstractmethod + async def rerank( + self, query: str, results: list[SearchResult], top_n: int + ) -> list[SearchResult]: + """Return the `top_n` most relevant candidates, most relevant first. + + The returned results carry the reranker's relevance score, not the + vector distance they arrived with, so a caller ordering or thresholding + on `score` reads the reranker's judgement. + """ + + +class CohereReranker(BaseReranker): + """Cohere's rerank endpoint behind :class:`BaseReranker`. + + The client is built on first use and can be injected, so a test drives the + reranker without a network or a key. A reranker is built per search and + reranks once, so a client it builds itself is request-scoped: `rerank` + closes it before returning rather than leaving the httpx connection pool + open. Nothing else does - the process-wide retrieval service builds a fresh + reranker every search - so an unclosed client would accumulate one pool per + query. An injected client belongs to its caller and is left open. + """ + + def __init__(self, model: str, api_key: str, client: AsyncClientV2 | None = None) -> None: + self.model = model + self._api_key = api_key + self._client = client + self._owns_client = client is None + + def __eq__(self, other: object) -> bool: + # Two rerankers are the same reranker when they name the same model and + # pay with the same key. Retrieval compares them across a multi-collection + # union to decide whether the bound collections share one configuration - + # only then may one reranker reorder the whole union on one credential. + if not isinstance(other, CohereReranker): + return NotImplemented + return self.model == other.model and self._api_key == other._api_key + + def __hash__(self) -> int: + return hash((self.model, self._api_key)) + + @property + def client(self) -> AsyncClientV2: + if self._client is None: + self._client = cohere.AsyncClientV2(api_key=self._api_key) + return self._client + + async def rerank( + self, query: str, results: list[SearchResult], top_n: int + ) -> list[SearchResult]: + if not results: + return [] + + client = self.client + try: + response = await client.rerank( + model=self.model, + query=query, + documents=[r.content for r in results], + top_n=min(top_n, len(results)), + ) + + # Booked only once the call has returned: Cohere does not bill a + # failed request, and a raise propagates to retrieval, which degrades + # to the un-reranked order rather than failing the search. + book_ambient_spend(self._spend_entry(response, len(results))) + + return [ + SearchResult( + content=results[item.index].content, + score=item.relevance_score, + metadata=results[item.index].metadata, + parent_doc_id=results[item.index].parent_doc_id, + ) + for item in response.results + ] + finally: + await self._release(client) + + async def _release(self, client: AsyncClientV2) -> None: + """Close a client this reranker built; leave an injected one alone. + + Best-effort: a failure to return the connection pool is logged, never + raised, so it cannot mask the rerank's own result or exception. + """ + if not self._owns_client: + return + self._client = None + try: + await client.__aexit__(None, None, None) + except Exception: + logger.warning("[RERANK] Closing the Cohere client failed", exc_info=True) + + def _spend_entry(self, response: object, candidate_count: int) -> SpendEntry: + return SpendEntry( + model_name=self.model, + input_tokens=0, + output_tokens=0, + cost_usd=_PRICE_PER_SEARCH_UNIT_USD * self._billed_units(response, candidate_count), + priced=True, + ) + + @staticmethod + def _billed_units(response: object, candidate_count: int) -> int: + """The search units Cohere actually billed, or an estimate from the count. + + Cohere splits a document past its token threshold into several billable + documents, so a request with fewer than `_DOCS_PER_SEARCH_UNIT` + candidates can still cost more than one unit - the ingestion config + allows chunks large enough for this to happen. The response carries the + real figure at `meta.billed_units.search_units`; every level of that + chain is optional, so when it is absent fall back to estimating one unit + per `_DOCS_PER_SEARCH_UNIT` candidates. + """ + meta = getattr(response, "meta", None) + billed = getattr(meta, "billed_units", None) + search_units = getattr(billed, "search_units", None) + if search_units is not None: + return ceil(search_units) + return ceil(candidate_count / _DOCS_PER_SEARCH_UNIT) + + +async def build_reranker( + collection_name: str, organization_id: UUID | None, knowledge_base_id: UUID | None = None +) -> BaseReranker | None: + """Bind a collection's resolved rerank credential to a concrete reranker. + + The one composition point for reranking: resolution answers whether a + collection is configured and with whose key, and this turns that into the + single implementation there is. Shared by every path that retrieves - the + `/rag/search` route and the agent-run knowledge tool alike - so reranking is + wired the same way in both, and a second provider is a branch here rather + than a change at each call site. + + `knowledge_base_id` pins resolution to the knowledge base the caller was + authorized against, rather than one looked up by the non-unique collection + name (#913); see `reranker_for_collection`. + """ + resolved = await reranker_for_collection(collection_name, organization_id, knowledge_base_id) + if resolved is None: + return None + return CohereReranker(model=resolved.model, api_key=resolved.api_key) diff --git a/backend/app/services/rag/retrieval.py b/backend/app/services/rag/retrieval.py index 1ff0685c0..97fdb3401 100644 --- a/backend/app/services/rag/retrieval.py +++ b/backend/app/services/rag/retrieval.py @@ -4,15 +4,30 @@ import logging import time from abc import ABC, abstractmethod +from collections.abc import Awaitable, Callable +from uuid import UUID from rank_bm25 import BM25Okapi from app.services.rag.config import RAGSettings from app.services.rag.models import SearchResult +from app.services.rag.reranker import BaseReranker from app.services.rag.vectorstore import BaseVectorStore logger = logging.getLogger(__name__) +# How a retrieval service learns whether a collection reranks, and with what. +# Async because the answer lives in the database, injected so the store never +# imports platform policy - the same shape as the embedding resolver. +RerankerResolver = Callable[[str, UUID | None, UUID | None], Awaitable[BaseReranker | None]] + +# Recall overfetches so min-score filtering and dedup still leave `limit` +# results. A reranker wants a wider net than that - the point of it is to +# surface a good answer sitting well below the top by distance - so it fetches +# more and truncates after reordering. +_DEFAULT_FETCH_MULTIPLIER = 2 +_RERANK_FETCH_MULTIPLIER = 4 + def _result_key(r: SearchResult) -> str: if r.parent_doc_id: @@ -29,6 +44,9 @@ async def retrieve( limit: int = 5, min_score: float = 0.0, filter: str = "", + *, + organization_id: UUID | None, + knowledge_base_id: UUID | None = None, ) -> list[SearchResult]: pass @@ -38,10 +56,14 @@ def __init__( self, vector_store: BaseVectorStore, settings: RAGSettings, + reranker_resolver: RerankerResolver | None = None, ): self.store = vector_store self.settings = settings self._hybrid_enabled = settings.enable_hybrid_search + # None leaves retrieval byte-for-byte its pre-reranker self: every path + # resolves no reranker and truncates by distance, exactly as before. + self._reranker_resolver = reranker_resolver @staticmethod def _rrf_fuse( @@ -76,14 +98,23 @@ def _rrf_fuse( ] async def _bm25_search( - self, query: str, collection_name: str, limit: int + self, + query: str, + collection_name: str, + limit: int, + organization_id: UUID | None, + knowledge_base_id: UUID | None = None, ) -> list[SearchResult]: docs = await self.store.get_documents(collection_name) if not docs: return [] all_results = await self.store.search( - collection_name=collection_name, query=query, limit=min(limit * 10, 100) + collection_name=collection_name, + query=query, + limit=min(limit * 10, 100), + organization_id=organization_id, + knowledge_base_id=knowledge_base_id, ) if not all_results: return [] @@ -105,6 +136,65 @@ async def _bm25_search( if s > 0 ] + async def _reranker_for( + self, + collection_name: str, + organization_id: UUID | None, + knowledge_base_id: UUID | None = None, + ) -> BaseReranker | None: + """The reranker one collection uses, or None when none is configured. + + None whenever no resolver was injected, so a service built without one + never reranks and never touches the database looking for a key. + + `knowledge_base_id`, when the caller has one, pins resolution to the + knowledge base access already authorized rather than one looked up by + the non-unique collection name; `organization_id` is the fallback scope + for callers with no authorized identity of their own (#913). + """ + if self._reranker_resolver is None: + return None + return await self._reranker_resolver(collection_name, organization_id, knowledge_base_id) + + @staticmethod + def _shared_reranker(rerankers: list[BaseReranker | None]) -> BaseReranker | None: + """The one reranker every collection agrees on, or None if they do not. + + None the moment any collection resolves to a different reranker or to + none at all: a union is reranked only when all of its collections share + one, so a disabled or differently-keyed collection in the set turns + reranking off for the whole union rather than reordering it on a + credential that is not its own. + """ + first = rerankers[0] if rerankers else None + if first is None: + return None + return first if all(r == first for r in rerankers) else None + + @staticmethod + async def _rank_and_truncate( + reranker: BaseReranker | None, + query: str, + candidates: list[SearchResult], + limit: int, + ) -> list[SearchResult]: + """Rerank the candidates and keep the top `limit`, or just keep the top. + + A reranker failure degrades to the by-distance order rather than failing + the search: reranking is an improvement on a working retrieval, and a + Cohere outage must not take knowledge search down with it. The + misconfiguration cases never reach here - resolution already turned + those into no reranker at all - so a raise here is a runtime fault worth + a log line. + """ + if reranker is None: + return candidates[:limit] + try: + return await reranker.rerank(query, candidates, limit) + except Exception: + logger.warning("[RETRIEVAL] Reranking failed; falling back to distance order") + return candidates[:limit] + async def retrieve( self, query: str, @@ -112,10 +202,49 @@ async def retrieve( limit: int = 5, min_score: float = 0.0, filter: str = "", + *, + organization_id: UUID | None, + knowledge_base_id: UUID | None = None, ) -> list[SearchResult]: - # Overfetch so min-score filtering and dedup still leave `limit` results. - fetch_multiplier = 2 + reranker = await self._reranker_for(collection_name, organization_id, knowledge_base_id) + multiplier = _RERANK_FETCH_MULTIPLIER if reranker else _DEFAULT_FETCH_MULTIPLIER + candidates = await self._recall( + query, + collection_name, + limit, + min_score, + filter, + fetch_multiplier=multiplier, + organization_id=organization_id, + knowledge_base_id=knowledge_base_id, + ) + return await self._rank_and_truncate(reranker, query, candidates, limit) + async def _recall( + self, + query: str, + collection_name: str, + limit: int, + min_score: float, + filter: str, + *, + fetch_multiplier: int, + organization_id: UUID | None, + knowledge_base_id: UUID | None = None, + ) -> list[SearchResult]: + """Vector (and optionally BM25) recall, filtered and deduplicated. + + Everything retrieval does before ranking: the candidate set, tagged with + the collection each result came from, not yet truncated to `limit`. Held + apart from `retrieve` so a multi-collection search can gather candidates + from several collections and rerank the union once, rather than reranking + each collection and merging the winners. + + `min_score` gates this recall on the vector-distance score the store + returns. It is not re-applied after reranking, where `score` carries the + reranker's relevance judgement on a different scale - so a caller must not + threshold on a reranked result's `score` as if it were the recall score. + """ logger.info( "[RETRIEVAL] Query: '%.50s...', collection: %s, limit: %d, filter: '%s'", query, @@ -131,6 +260,8 @@ async def retrieve( query=query, filter_expr=filter, limit=limit * fetch_multiplier, + organization_id=organization_id, + knowledge_base_id=knowledge_base_id, ) search_time = time.time() - start_time @@ -141,7 +272,9 @@ async def retrieve( ) if self._hybrid_enabled: - bm25_results = await self._bm25_search(query, collection_name, limit * fetch_multiplier) + bm25_results = await self._bm25_search( + query, collection_name, limit * fetch_multiplier, organization_id, knowledge_base_id + ) if bm25_results: pipeline_results = self._rrf_fuse(pipeline_results, bm25_results) logger.info("[RETRIEVAL] Hybrid search: fused %d results", len(pipeline_results)) @@ -179,23 +312,23 @@ async def retrieve( r.content, ) - final_results = deduped_results[:limit] - # Which collection answered, on every result rather than only when several # were searched. A caller cannot derive it - one search may span bases and # two bases may share a collection - and a chunk whose origin is unknown - # cannot be cited, which is the whole job of a retrieval result. - for r in final_results: + # cannot be cited, which is the whole job of a retrieval result. Stamped + # here on every candidate so it survives reranking, which builds fresh + # results carrying this metadata forward. + for r in deduped_results: r.metadata["collection"] = collection_name total_time = time.time() - start_time logger.info( - "[RETRIEVAL] Total retrieval time: %.3fs, returning %d results", + "[RETRIEVAL] Total recall time: %.3fs, %d candidates", total_time, - len(final_results), + len(deduped_results), ) - return final_results + return deduped_results async def retrieve_multi( self, @@ -203,6 +336,9 @@ async def retrieve_multi( collection_names: list[str], limit: int = 5, min_score: float = 0.0, + *, + organization_id: UUID | None, + knowledge_base_ids: list[UUID] | None = None, ) -> list[SearchResult]: """Search several collections and merge what they return. @@ -214,17 +350,54 @@ async def retrieve_multi( A collection nobody has ingested into is not a failure: its table does not exist yet, and the store reports that as no results. + + When a reranker is configured it runs once over the union of every + collection's candidates, not per collection: reranking each collection + separately then merging the winners would rank against the wrong pool. + + One reranker reorders the union only when *every* collection resolves to + the same one - same model, same key. On the agent-run path that always + holds: an agent's bound collections share one organization and one + configuration. But `/rag/search` may pass any readable set of one + organization, and a set whose collections disagree - one reranking, one + not, or two on different keys - is not reranked at all: reranking a + disabled collection's candidates on another collection's credential would + send content that opted out to Cohere and bill it to the wrong key. A + mixed set falls back to the by-distance union rather than picking a winner + by position. Absent a shared reranker this is byte-for-byte the previous + merge - each collection's top `limit`, fused, sorted, deduplicated, + truncated. """ + # Each collection carries the id of the knowledge base access authorized + # for it, so resolution reads that row rather than one looked up by the + # non-unique name (#913). The lists are built together at the one call + # site, so they stay aligned; a caller with no authorized identity (an + # agent run, the CLI) passes none and every collection falls back to the + # organization-scoped lookup. + kb_ids: list[UUID | None] = [ + knowledge_base_ids[i] if knowledge_base_ids is not None else None + for i in range(len(collection_names)) + ] + rerankers = [ + await self._reranker_for(name, organization_id, kb_id) + for name, kb_id in zip(collection_names, kb_ids) + ] + reranker = self._shared_reranker(rerankers) + multiplier = _RERANK_FETCH_MULTIPLIER if reranker else _DEFAULT_FETCH_MULTIPLIER + all_results: list[SearchResult] = [] - for name in collection_names: - all_results.extend( - await self.retrieve( - query=query, - collection_name=name, - limit=limit, - min_score=min_score, - ) + for name, kb_id in zip(collection_names, kb_ids): + recalled = await self._recall( + query, + name, + limit, + min_score, + "", + fetch_multiplier=multiplier, + organization_id=organization_id, + knowledge_base_id=kb_id, ) + all_results.extend(recalled if reranker else recalled[:limit]) all_results.sort(key=lambda r: r.score, reverse=True) @@ -236,4 +409,4 @@ async def retrieve_multi( seen_keys.add(key) deduped.append(r) - return deduped[:limit] + return await self._rank_and_truncate(reranker, query, deduped, limit) diff --git a/backend/app/services/rag/vectorstore.py b/backend/app/services/rag/vectorstore.py index 08adecfad..51c49898c 100644 --- a/backend/app/services/rag/vectorstore.py +++ b/backend/app/services/rag/vectorstore.py @@ -3,6 +3,7 @@ from abc import ABC, abstractmethod from pathlib import Path from typing import Any +from uuid import UUID # Registers every model table on `Base.metadata`, which `list_collections` judges a # `rag_` table against and `_table` refuses a collection name against. Another import @@ -57,12 +58,21 @@ def _document_is_unaddressed(doc: DocumentInfo) -> bool: class BaseVectorStore(ABC): @abstractmethod - async def insert_document(self, collection_name: str, document: Document) -> None: + async def insert_document( + self, collection_name: str, document: Document, *, organization_id: UUID | None + ) -> None: pass @abstractmethod async def search( - self, collection_name: str, query: str, limit: int = 4, filter_expr: str = "" + self, + collection_name: str, + query: str, + limit: int = 4, + filter_expr: str = "", + *, + organization_id: UUID | None, + knowledge_base_id: UUID | None = None, ) -> list[SearchResult]: pass @@ -75,7 +85,9 @@ async def delete_document(self, collection_name: str, document_id: str) -> None: pass @abstractmethod - async def get_collection_info(self, collection_name: str) -> CollectionInfo: + async def get_collection_info( + self, collection_name: str, *, organization_id: UUID | None + ) -> CollectionInfo: pass @abstractmethod @@ -156,7 +168,7 @@ async def find_existing_document( by_hash = doc return by_filename or by_hash - async def create_collection(self, name: str) -> None: + async def create_collection(self, name: str, *, organization_id: UUID | None) -> None: """Make the collection's backing objects, refusing a name that cannot have any. The check is here as well as in `_table` because a subclass is free to @@ -169,7 +181,7 @@ async def create_collection(self, name: str) -> None: :func:`app.db.vector_tables.validate_collection_name`. """ validate_collection_name(name, metadata=Base.metadata) - await self._ensure_collection(name) + await self._ensure_collection(name, organization_id) def _build_chunk_metadata( self, chunk: "DocumentPageChunk", document: Document @@ -234,7 +246,7 @@ def _group_documents(self, results: list[dict[str, Any]]) -> list[DocumentInfo]: # How a store learns which model a collection embeds with. Async because the # answer lives in the database, injected so the template's store never imports # platform policy. -EmbeddingResolver = Callable[[str], Awaitable[ResolvedEmbeddings | None]] +EmbeddingResolver = Callable[[str, UUID | None, UUID | None], Awaitable[ResolvedEmbeddings | None]] # pgvector's HNSW builds over a `vector` column only up to this width; past it, # `CREATE INDEX` fails with "column cannot have more than 2000 dimensions for @@ -329,7 +341,9 @@ def _table(self, name: str) -> str: validate_collection_name(name, metadata=Base.metadata) return f"{VECTOR_TABLE_PREFIX}{name}" - async def _for_collection(self, name: str) -> tuple[EmbeddingService, int]: + async def _for_collection( + self, name: str, organization_id: UUID | None, knowledge_base_id: UUID | None = None + ) -> tuple[EmbeddingService, int]: """The embedder and vector width this one collection uses. Cached per (collection, model, key): an `EmbeddingService` holds an @@ -345,7 +359,7 @@ async def _for_collection(self, name: str) -> tuple[EmbeddingService, int]: The recorded width wins over the catalog's: the table was created at that number. """ - resolved = await self._resolver(name) + resolved = await self._resolver(name, organization_id, knowledge_base_id) if resolved is None: return self.embedder, self.dim cache_key = (name, resolved.model, resolved.api_key, resolved.base_url) @@ -383,10 +397,10 @@ def _distance_expr(dim: int) -> str: return f"(embedding::halfvec({dim}))" return "embedding" - async def _ensure_collection(self, name: str) -> None: + async def _ensure_collection(self, name: str, organization_id: UUID | None) -> None: """Create table for collection if not exists.""" table = self._table(name) - _, dim = await self._for_collection(name) + _, dim = await self._for_collection(name, organization_id) operator_class = "halfvec_cosine_ops" if dim > _HNSW_MAX_VECTOR_DIM else "vector_cosine_ops" async with self.async_session() as session: await session.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) @@ -444,7 +458,9 @@ async def _collection_exists(self, name: str) -> bool: ) return result.scalar() is not None - async def insert_document(self, collection_name: str, document: Document) -> None: + async def insert_document( + self, collection_name: str, document: Document, *, organization_id: UUID | None + ) -> None: """Write a document's chunks, a batch of rows per statement. One statement per `_CHUNK_INSERT_BATCH` chunks rather than one per chunk: @@ -464,10 +480,10 @@ async def insert_document(self, collection_name: str, document: Document) -> Non while leaving the worker's memory exactly where it was. """ table = self._table(collection_name) - await self._ensure_collection(collection_name) + await self._ensure_collection(collection_name, organization_id) if not document.chunked_pages: raise ValueError("Document has no chunked pages.") - embedder, _ = await self._for_collection(collection_name) + embedder, _ = await self._for_collection(collection_name, organization_id) vectors = embedder.embed_document(document) statement = text(f""" INSERT INTO {table} (id, parent_doc_id, content, embedding, metadata) @@ -492,7 +508,14 @@ async def insert_document(self, collection_name: str, document: Document) -> Non await session.commit() async def search( - self, collection_name: str, query: str, limit: int = 4, filter_expr: str = "" + self, + collection_name: str, + query: str, + limit: int = 4, + filter_expr: str = "", + *, + organization_id: UUID | None, + knowledge_base_id: UUID | None = None, ) -> list[SearchResult]: """Nearest chunks in a collection, reporting an absent one as empty. @@ -502,11 +525,17 @@ async def search( knowledge base nobody has uploaded to yet turned asyncpg's `UndefinedTableError` into a 500, and it is checked before embedding so an empty collection costs no embedding call either. + + `knowledge_base_id` pins the query's embedding config to the knowledge + base the caller was authorized against rather than one looked up by the + non-unique collection name (#913). """ table = self._table(collection_name) if not await self._collection_exists(collection_name): return [] - embedder, dim = await self._for_collection(collection_name) + embedder, dim = await self._for_collection( + collection_name, organization_id, knowledge_base_id + ) query_vector = embedder.embed_query(query) # Parse the shared `parent_doc_id == ""` filter format and apply @@ -550,7 +579,9 @@ async def search( for row in rows ] - async def get_collection_info(self, collection_name: str) -> CollectionInfo: + async def get_collection_info( + self, collection_name: str, *, organization_id: UUID | None + ) -> CollectionInfo: """Vector count for a collection, reporting an absent one as empty. A collection's table is created lazily by the first ingest, so "no table" @@ -562,7 +593,7 @@ async def get_collection_info(self, collection_name: str) -> CollectionInfo: question with an empty list; its comment claimed this method already did the same, and now it does. """ - _, dim = await self._for_collection(collection_name) + _, dim = await self._for_collection(collection_name, organization_id) if not await self._collection_exists(collection_name): return CollectionInfo(name=collection_name, total_vectors=0, dim=dim) table = self._table(collection_name) diff --git a/backend/app/services/rag_document.py b/backend/app/services/rag_document.py index c62e0e431..af6bbd39c 100644 --- a/backend/app/services/rag_document.py +++ b/backend/app/services/rag_document.py @@ -255,7 +255,7 @@ async def dispatch_upload( ) doc_id = rag_doc.id - await vector_store.create_collection(collection_name) + await vector_store.create_collection(collection_name, organization_id=organization_id) await self._queue_parse( doc_id, diff --git a/backend/app/services/rerank_resolution.py b/backend/app/services/rerank_resolution.py new file mode 100644 index 000000000..d77286e1d --- /dev/null +++ b/backend/app/services/rerank_resolution.py @@ -0,0 +1,178 @@ +"""Whether a collection reranks its search results - and whose key pays. + +The sibling of :mod:`app.services.embedding_resolution`, and deliberately its +mirror image in shape. Retrieval asks per collection whether a reranker is +configured; the answer carries the reranker's model and the organization key it +runs on, or nothing. + +Where embeddings *fall back* to the deployment key when a collection's chosen +one is gone, reranking *turns off*. There is no deployment reranker key - the +feature is off by default and on only when a collection names both a model and a +usable organization secret - so every path that is not "a usable key the +collection chose" resolves to `None` and retrieval behaves exactly as it did +before the feature. The distinction still has to be *said*, though: a collection +that chose a key and lost it is a misconfiguration an operator should see, where +a collection that chose nothing is the normal off state and must stay quiet. So +resolution classifies the reason with :class:`RerankKeySource` and logs the +three degraded ones, exactly as embedding resolution names its own. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from enum import StrEnum +from uuid import UUID + +from sqlalchemy.ext.asyncio import AsyncSession + +from app.core.secret_kinds import ApiKeySecret, SecretKind, unseal_secret +from app.core.vault import VaultScope +from app.db.models.knowledge_base import KnowledgeBase +from app.db.session import get_db_context +from app.repositories import knowledge_base_repo, organization_secret_repo + +logger = logging.getLogger(__name__) + +# Secret purposes that can pay for reranking. Cohere is the only reranker today, +# and `cohere` is already a model-provider purpose - one Cohere API key reranks +# and chats alike - so this reuses it rather than minting a second entry that +# would collide with it. The tuple exists so a second provider is one entry. +RERANK_KEY_PURPOSES = ("cohere",) + +# The rerank models this platform can actually run. Cohere Rerank 3.5 is the only +# one behind `CohereReranker` today; a value outside this set would be accepted, +# stored and displayed as configured, then fail every search inside Cohere - a +# request `_rank_and_truncate` swallows, so reranking would be silently off. So a +# model is validated at create and update against this tuple, which grows by one +# entry when a second model is supported. +SUPPORTED_RERANK_MODELS = ("rerank-v3.5",) + + +class RerankKeySource(StrEnum): + """Why a collection did or did not get a reranker. + + Only :attr:`CONFIGURED` yields one. The rest all mean "no reranking", and + they are kept apart for the same reason embedding resolution keeps its + sources apart: telling "the collection chose no reranker" from "the key it + chose is gone" is the difference between silence and a line an operator + needs to see. + """ + + CONFIGURED = "configured" + NOT_CONFIGURED = "not_configured" + SECRET_MISSING = "secret_missing" + SECRET_UNUSABLE = "secret_unusable" + SECRET_WRONG_KIND = "secret_wrong_kind" + + @property + def is_degraded(self) -> bool: + """True when the collection asked for a reranker and did not get one. + + `NOT_CONFIGURED` is not degraded: a collection that named no reranker is + supposed to have none, and warning on every unconfigured search would + bury the three that mean a real misconfiguration. + """ + return self in _DEGRADED + + +_DEGRADED = frozenset( + { + RerankKeySource.SECRET_MISSING, + RerankKeySource.SECRET_UNUSABLE, + RerankKeySource.SECRET_WRONG_KIND, + } +) + + +@dataclass(frozen=True, repr=False) +class ResolvedReranker: + """Everything one collection's rerank call needs. + + `repr=False` for the same reason :class:`ResolvedEmbeddings` carries it: the + dataclass holds a plaintext key, and the default repr is how a key reaches a + log line. + """ + + model: str + api_key: str + + def __repr__(self) -> str: + return f"ResolvedReranker(model={self.model!r}, api_key='***')" + + +async def reranker_for_collection( + collection_name: str, organization_id: UUID | None, knowledge_base_id: UUID | None = None +) -> ResolvedReranker | None: + """Resolve one collection's reranker, or `None` if it has none. + + `None` for a collection no knowledge base claims, for one that named no + reranker, and for one whose chosen key is missing, unusable or the wrong + kind - the last three with a warning, because they are a misconfiguration + rather than the off state. Opens its own session because retrieval reaches + here from places with no request in sight: an agent mid-run, a direct search. + + `knowledge_base_id`, when given, is the knowledge base the caller was already + authorized against, and resolution reads *that* row rather than looking one + up by name. `collection_name` is not unique, and the access check and a + name+organization lookup can pick different rows for the same name - an `app` + collection everyone may read and a restricted `org` collection of the same + name - so resolving by name could unseal and spend the key of a row the + caller was never granted (#913). The search path passes the authorized id; + ingestion and the CLI, which choose the row themselves and have no distinct + authorized identity, pass none and fall back to the `organization_id`-scoped + lookup. + """ + async with get_db_context() as db: + kb = ( + await knowledge_base_repo.get_by_id(db, knowledge_base_id) + if knowledge_base_id is not None + else await knowledge_base_repo.get_for_collection(db, collection_name, organization_id) + ) + if kb is None: + return None + resolved, source = await _resolve_reranker(db, kb) + if source.is_degraded: + logger.warning("rerank_%s", source.value, extra={"collection": collection_name}) + return resolved + + +async def _resolve_reranker( + db: AsyncSession, kb: KnowledgeBase +) -> tuple[ResolvedReranker | None, RerankKeySource]: + """The reranker a collection is configured for, or `None` and why not. + + Unlike embedding resolution, no failure lands on a deployment key: there is + none, so every path but a usable organization secret returns `None`. The + second element is what stops that from being invisible - it is carried out to + the log line above. + """ + model = kb.rerank_model + secret_id = kb.rerank_secret_id + organization_id = kb.organization_id + if model is None and secret_id is None: + return None, RerankKeySource.NOT_CONFIGURED + if model is None or secret_id is None or organization_id is None: + # A half-configured reranker, not the off state. The pair is written + # together - create and update enforce it - but deleting the chosen + # secret nulls rerank_secret_id through the foreign key while leaving + # rerank_model set, and that stopped reranking with no signal at all + # until this told the half state apart from the null/null off state. + return None, RerankKeySource.SECRET_MISSING + + row = await organization_secret_repo.get(db, secret_id, organization_id=organization_id) + if row is None: + return None, RerankKeySource.SECRET_MISSING + try: + secret = unseal_secret( + row.sealed_secret, + kind=SecretKind(row.kind), + scope=VaultScope.organization(organization_id), + key_version=row.key_version, + ) + except Exception: + return None, RerankKeySource.SECRET_UNUSABLE + if not isinstance(secret, ApiKeySecret): + return None, RerankKeySource.SECRET_WRONG_KIND + resolved = ResolvedReranker(model=model, api_key=secret.api_key.get_secret_value()) + return resolved, RerankKeySource.CONFIGURED diff --git a/backend/app/services/stats.py b/backend/app/services/stats.py index 6b859db07..082705913 100644 --- a/backend/app/services/stats.py +++ b/backend/app/services/stats.py @@ -25,6 +25,7 @@ from app.core.exceptions import AuthorizationError, ValidationError from app.core.permissions import Perm +from app.db.models.ingestion_spend import SpendSource from app.repositories import ( agent_run_repo, ingestion_spend_repo, @@ -272,20 +273,52 @@ async def usage( previous_model_usd = await agent_run_repo.sum_cost_window( self.db, organization_id=org, start=prev_start, end=prev_end, where=where ) + # Indexing and search are both organization-wide non-run spend, kept + # apart so a search is not reported as indexing (they share the table, + # `source` tells them apart). Both are unnarrowable - a worker's sync + # and a colleague's search belong to no one person or agent - so a + # narrowed window reports model spend alone. ingestion_usd = Decimal(0) previous_ingestion_usd = Decimal(0) + retrieval_usd = Decimal(0) + previous_retrieval_usd = Decimal(0) if where == RunFilter(): ingestion_usd = await ingestion_spend_repo.sum_cost_window( - self.db, organization_id=org, start=window.start, end=window.end + self.db, + organization_id=org, + start=window.start, + end=window.end, + source=SpendSource.INGESTION, ) previous_ingestion_usd = await ingestion_spend_repo.sum_cost_window( - self.db, organization_id=org, start=prev_start, end=prev_end + self.db, + organization_id=org, + start=prev_start, + end=prev_end, + source=SpendSource.INGESTION, + ) + retrieval_usd = await ingestion_spend_repo.sum_cost_window( + self.db, + organization_id=org, + start=window.start, + end=window.end, + source=SpendSource.RETRIEVAL, + ) + previous_retrieval_usd = await ingestion_spend_repo.sum_cost_window( + self.db, + organization_id=org, + start=prev_start, + end=prev_end, + source=SpendSource.RETRIEVAL, ) cost = CostBlock( - period_usd=model_usd + ingestion_usd, - previous_period_usd=previous_model_usd + previous_ingestion_usd, + period_usd=model_usd + ingestion_usd + retrieval_usd, + previous_period_usd=( + previous_model_usd + previous_ingestion_usd + previous_retrieval_usd + ), model_usd=model_usd, ingestion_usd=ingestion_usd, + retrieval_usd=retrieval_usd, by_provider=[ ProviderCost(provider=provider, cost_usd=cost_usd) for provider, cost_usd in await agent_run_repo.cost_by_provider_window( diff --git a/backend/app/worker/tasks/rag_tasks.py b/backend/app/worker/tasks/rag_tasks.py index f788d97c5..6d2341b99 100644 --- a/backend/app/worker/tasks/rag_tasks.py +++ b/backend/app/worker/tasks/rag_tasks.py @@ -98,8 +98,14 @@ def _announcing_resolver() -> EmbeddingResolver: """ announced: set[str] = set() - async def resolve(collection_name: str) -> ResolvedEmbeddings | None: - resolved = await embeddings_for_collection(collection_name) + async def resolve( + collection_name: str, + organization_id: UUID | None, + knowledge_base_id: UUID | None = None, + ) -> ResolvedEmbeddings | None: + resolved = await embeddings_for_collection( + collection_name, organization_id, knowledge_base_id + ) if ( resolved is not None and resolved.key_source.is_degraded @@ -113,7 +119,9 @@ async def resolve(collection_name: str) -> ResolvedEmbeddings | None: @asynccontextmanager -async def _ingestion_service(*, processor: DocumentProcessor) -> AsyncIterator[IngestionService]: +async def _ingestion_service( + *, processor: DocumentProcessor, organization_id: UUID | None +) -> AsyncIterator[IngestionService]: """An ingester that reads documents the way the collection asked to be read. Both halves come off the collection. The parser, the chunker and the image @@ -152,6 +160,7 @@ async def _ingestion_service(*, processor: DocumentProcessor) -> AsyncIterator[I resolver=_announcing_resolver(), engine=engine, ), + organization_id=organization_id, ) finally: await engine.dispose() @@ -210,11 +219,7 @@ async def _knowledge_base_for( """ if collection_name is None: return None - candidates = await knowledge_base_repo.list_by_collection_name(db, collection_name) - for kb in candidates: - if organization_id is None or kb.organization_id == organization_id: - return kb - return next((kb for kb in candidates if kb.organization_id is None), None) + return await knowledge_base_repo.get_for_collection(db, collection_name, organization_id) async def _config_for_collection( @@ -332,7 +337,7 @@ async def _run_ingestion( ledger = SpendLedger(organization_id=organization_id) file_path = Path(filepath) - async with _ingestion_service(processor=processor) as ingester: + async with _ingestion_service(processor=processor, organization_id=organization_id) as ingester: try: with metered_by(ledger): result = await ingester.ingest_file( @@ -424,7 +429,7 @@ async def _run_sync( # Entered after the validations above, so an early "path not found" return # builds no engine, and every return inside the loop still disposes one (#948). - async with _ingestion_service(processor=processor) as ingester: + async with _ingestion_service(processor=processor, organization_id=None) as ingester: for filepath in files: async with get_worker_db_context() as db: sync_log_check = await RAGSyncService(db).get_sync_log(sync_log_id) @@ -776,7 +781,7 @@ async def _run_source_sync(source_id: str, sync_log_id: str | None = None) -> di total = 0 ledger = SpendLedger(organization_id=organization_id) - async with _ingestion_service(processor=processor) as ingester: + async with _ingestion_service(processor=processor, organization_id=organization_id) as ingester: try: files = await connector.list_files(config, credential) total = len(files) diff --git a/backend/pyproject.toml b/backend/pyproject.toml index d8187eb17..2c536e15a 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -541,6 +541,7 @@ include = [ "app/services/collection_access.py", "app/services/embed_session.py", "app/services/embedding_resolution.py", + "app/services/rerank_resolution.py", "app/services/approvals.py", "app/services/audit.py", "app/services/health.py", @@ -552,6 +553,7 @@ include = [ # mention, a schedule — so a silent failure here is indistinguishable from # an agent that simply had nothing to say. "app/services/notifications.py", + "app/services/knowledge_search.py", "app/services/mcp_catalog.py", "app/services/portal_catalog.py", "app/services/portals/__init__.py", @@ -778,6 +780,7 @@ include = [ # ungated, which is the failure this gate exists to name (#663). "app/services/embed_session.py", "app/services/embedding_resolution.py", + "app/services/rerank_resolution.py", "app/services/approvals.py", # The read half of the audit trail. `audit:read` is the only way to see who # did what in an organization, and the scope is the whole point of the @@ -810,6 +813,7 @@ include = [ # mention, a schedule — so a silent failure here is indistinguishable from # an agent that simply had nothing to say. "app/services/notifications.py", + "app/services/knowledge_search.py", "app/services/mcp_catalog.py", "app/services/portal_catalog.py", "app/services/portals/__init__.py", diff --git a/backend/tests/api/test_error_envelope.py b/backend/tests/api/test_error_envelope.py index 589054beb..42390bc30 100644 --- a/backend/tests/api/test_error_envelope.py +++ b/backend/tests/api/test_error_envelope.py @@ -495,7 +495,9 @@ async def test_a_failed_knowledge_search_names_the_collections_not_the_upstream( ), pytest.raises(ExternalServiceError) as refusal, ): - await search_knowledge_base(query="our refund policy", kb_collection_names=["kb_ops"]) + await search_knowledge_base( + query="our refund policy", kb_collection_names=["kb_ops"], organization_id=None + ) response = await self._refusal_on_the_wire(client, refusal.value) @@ -521,7 +523,9 @@ async def test_a_search_over_several_collections_names_the_operation_it_used( ), pytest.raises(ExternalServiceError) as refusal, ): - await search_knowledge_base(query="x", kb_collection_names=["kb_ops", "kb_hr"]) + await search_knowledge_base( + query="x", kb_collection_names=["kb_ops", "kb_hr"], organization_id=None + ) response = await self._refusal_on_the_wire(client, refusal.value) diff --git a/backend/tests/api/test_no_secret_escapes.py b/backend/tests/api/test_no_secret_escapes.py index 6dfdcfdec..d90d67e9f 100644 --- a/backend/tests/api/test_no_secret_escapes.py +++ b/backend/tests/api/test_no_secret_escapes.py @@ -111,6 +111,10 @@ "the same again: which vault key a collection embeds on - the id of a " "reference the organization can revoke, never the key itself" ), + "rerank_secret_id": ( + "and once more for reranking: which vault key a collection reranks on - " + "an id naming a revocable reference, never the Cohere key itself" + ), "token_secret_id": ( "the same, named for what it points at: an agent's Logfire write token " "lives in the vault and the spec carries only its id, because a spec is " diff --git a/backend/tests/integration/test_chunk_insert_batching.py b/backend/tests/integration/test_chunk_insert_batching.py index a04ae68e7..a69b38a0d 100644 --- a/backend/tests/integration/test_chunk_insert_batching.py +++ b/backend/tests/integration/test_chunk_insert_batching.py @@ -80,7 +80,7 @@ async def test_a_document_spanning_several_batches_writes_every_row( collection = f"batched_{uuid.uuid4().hex[:8]}" store = _store(engine) - await store.insert_document(collection, _document(chunks=250)) + await store.insert_document(collection, _document(chunks=250), organization_id=None) assert await _count(engine, f"rag_{collection}") == 250 @@ -95,10 +95,10 @@ async def test_re_inserting_the_same_chunks_updates_rather_than_duplicates( store = _store(engine) document = _document(chunks=10) - await store.insert_document(collection, document) + await store.insert_document(collection, document, organization_id=None) for chunk in document.chunked_pages: chunk.chunk_content = f"revised {chunk.chunk_num}" - await store.insert_document(collection, document) + await store.insert_document(collection, document, organization_id=None) table = f"rag_{collection}" assert await _count(engine, table) == 10 @@ -120,7 +120,7 @@ async def test_the_chunks_are_readable_back_in_document_order( document = _document(chunks=8) parent = document.chunked_pages[0].parent_doc_id - await store.insert_document(collection, document) + await store.insert_document(collection, document, organization_id=None) chunks = await store.get_document_chunks(collection, parent) assert [chunk.content for chunk in chunks] == [f"chunk {index}" for index in range(8)] diff --git a/backend/tests/integration/test_collection_name_tenant_isolation.py b/backend/tests/integration/test_collection_name_tenant_isolation.py new file mode 100644 index 000000000..cf4daeb73 --- /dev/null +++ b/backend/tests/integration/test_collection_name_tenant_isolation.py @@ -0,0 +1,211 @@ +"""Two organizations sharing a collection name each resolve their own config. + +`collection_name` is indexed but not unique, so two tenants can name a +collection the same thing. Resolving one by name alone returned whichever row +the database yielded first, which could unseal and bill another organization's +key (#913). The resolvers now take the organization the search or ingest acts +for; these run them against a real database with two tenants on one name and +assert each gets its own embedding and rerank configuration - never the other's. +""" + +from __future__ import annotations + +import uuid + +import pytest +from pydantic import SecretStr + +from app.core.secret_kinds import ApiKeySecret, seal_secret +from app.core.vault import VaultScope +from app.db.models.knowledge_base import KBScope, KnowledgeBase +from app.db.models.organization import Organization +from app.db.models.organization_secret import OrganizationSecret +from app.db.models.user import User +from app.services.embedding_resolution import embeddings_for_collection +from app.services.rerank_resolution import reranker_for_collection + +pytestmark = pytest.mark.anyio + +_SHARED = "shared_collection" + + +async def _org(db, name: str) -> Organization: + founder = User( + id=uuid.uuid4(), + email=f"{uuid.uuid4().hex}@example.com", + hashed_password="x", + is_active=True, + ) + db.add(founder) + await db.flush() + org = Organization( + id=uuid.uuid4(), + name=name, + slug=f"{name}-{uuid.uuid4().hex[:8]}", + created_by_user_id=founder.id, + ) + db.add(org) + await db.flush() + return org + + +async def _cohere_secret(db, org: Organization, key: str) -> OrganizationSecret: + sealed = seal_secret( + ApiKeySecret(api_key=SecretStr(key)), scope=VaultScope.organization(org.id) + ) + secret = OrganizationSecret( + id=uuid.uuid4(), + organization_id=org.id, + name="cohere", + kind="api_key", + purpose="cohere", + sealed_secret=sealed.ciphertext, + hint=sealed.hint, + key_version=sealed.key_version, + ) + db.add(secret) + await db.flush() + return secret + + +async def _kb(db, org: Organization, *, embedding_model: str, rerank_key: str) -> None: + secret = await _cohere_secret(db, org, rerank_key) + db.add( + KnowledgeBase( + id=uuid.uuid4(), + name=f"{org.name} handbook", + scope=KBScope.ORG.value, + collection_name=_SHARED, + embedding_model=embedding_model, + embedding_dim=1536, + rerank_model="rerank-v3.5", + rerank_secret_id=secret.id, + organization_id=org.id, + ingestion_config={}, + ) + ) + await db.flush() + + +async def test_each_organization_resolves_its_own_embedding_and_rerank_config(db) -> None: + org_a = await _org(db, "acme") + org_b = await _org(db, "globex") + await _kb(db, org_a, embedding_model="model-a", rerank_key="cohere-key-a") + await _kb(db, org_b, embedding_model="model-b", rerank_key="cohere-key-b") + # The resolvers open their own session, so the rows must be committed to be + # visible to it. + await db.commit() + + emb_a = await embeddings_for_collection(_SHARED, organization_id=org_a.id) + emb_b = await embeddings_for_collection(_SHARED, organization_id=org_b.id) + assert emb_a is not None and emb_a.model == "model-a" + assert emb_b is not None and emb_b.model == "model-b" + + rer_a = await reranker_for_collection(_SHARED, organization_id=org_a.id) + rer_b = await reranker_for_collection(_SHARED, organization_id=org_b.id) + assert rer_a is not None and rer_a.api_key == "cohere-key-a" + assert rer_b is not None and rer_b.api_key == "cohere-key-b" + + +async def test_an_organization_without_a_row_for_the_name_resolves_nothing(db) -> None: + """A third organization sharing neither row gets no config - not another + tenant's - so it can never unseal a key that is not its own (#913).""" + org_a = await _org(db, "acme") + await _kb(db, org_a, embedding_model="model-a", rerank_key="cohere-key-a") + stranger = await _org(db, "initech") + await db.commit() + + assert await embeddings_for_collection(_SHARED, organization_id=stranger.id) is None + assert await reranker_for_collection(_SHARED, organization_id=stranger.id) is None + + +def _kb_row(org: Organization, name: str, **secret_ids: uuid.UUID) -> KnowledgeBase: + return KnowledgeBase( + id=uuid.uuid4(), + name=name, + scope=KBScope.ORG.value, + collection_name=name, + embedding_model="model", + embedding_dim=1536, + organization_id=org.id, + ingestion_config={}, + **secret_ids, + ) + + +async def test_an_authorized_kb_id_pins_resolution_to_that_row_not_a_same_named_one(db) -> None: + """The residual of #913: an `app` KB and a restricted `org` KB can share a + name, and a name+organization lookup returns the org row (own-org wins) even + when access only authorized the app row. Passing the authorized `kb.id` + resolves *that* row, so a member who may read the app collection never + unseals or spends the restricted org collection's key.""" + org = await _org(db, "acme") + org_secret = await _cohere_secret(db, org, "org-only-key") + + app_kb = KnowledgeBase( + id=uuid.uuid4(), + name="shared (app)", + scope=KBScope.APP.value, + collection_name=_SHARED, + embedding_model="app-model", + embedding_dim=1536, + organization_id=None, + ingestion_config={}, + ) + org_kb = KnowledgeBase( + id=uuid.uuid4(), + name="shared (org, restricted)", + scope=KBScope.ORG.value, + collection_name=_SHARED, + embedding_model="org-model", + embedding_dim=1536, + rerank_model="rerank-v3.5", + rerank_secret_id=org_secret.id, + organization_id=org.id, + ingestion_config={}, + ) + db.add_all([app_kb, org_kb]) + await db.flush() + await db.commit() + + # By name+org, own-org wins: the restricted org row, and its key. + assert (await embeddings_for_collection(_SHARED, organization_id=org.id)).model == "org-model" + assert await reranker_for_collection(_SHARED, organization_id=org.id) is not None + + # Pinned to the authorized app row: the app config, and no org key. + pinned_emb = await embeddings_for_collection( + _SHARED, organization_id=org.id, knowledge_base_id=app_kb.id + ) + assert pinned_emb is not None and pinned_emb.model == "app-model" + assert ( + await reranker_for_collection(_SHARED, organization_id=org.id, knowledge_base_id=app_kb.id) + is None + ) + + +async def test_knowledge_bases_using_finds_embedding_and_rerank_bindings(db) -> None: + """A key bound as either a KB embedding or rerank credential is reported, so + the vault does not call it unused and invite a deletion that SET NULL then + turns off. Scoped to the organization: another tenant's binding never shows.""" + from app.repositories import knowledge_base_repo + + org = await _org(db, "acme") + other = await _org(db, "globex") + secret = await _cohere_secret(db, org, "co-key") + other_secret = await _cohere_secret(db, other, "other-key") + + db.add_all( + [ + _kb_row(org, "reranked", rerank_secret_id=secret.id), + _kb_row(org, "embedded", embedding_secret_id=secret.id), + _kb_row(org, "unrelated"), + _kb_row(other, "other-tenant", rerank_secret_id=other_secret.id), + ] + ) + await db.flush() + + found = await knowledge_base_repo.knowledge_bases_using( + db, organization_id=org.id, secret_id=secret.id + ) + + assert {name for _id, name in found} == {"reranked", "embedded"} diff --git a/backend/tests/integration/test_platform_flows.py b/backend/tests/integration/test_platform_flows.py index 27e27a8e7..d948d77d6 100644 --- a/backend/tests/integration/test_platform_flows.py +++ b/backend/tests/integration/test_platform_flows.py @@ -1293,7 +1293,7 @@ class _AcceptingStore: def __init__(self) -> None: self.created: list[str] = [] - async def create_collection(self, name: str) -> None: + async def create_collection(self, name: str, *, organization_id: object = None) -> None: self.created.append(name) diff --git a/backend/tests/integration/test_rag_existence_index.py b/backend/tests/integration/test_rag_existence_index.py index b9c31e6f7..44f93e1fa 100644 --- a/backend/tests/integration/test_rag_existence_index.py +++ b/backend/tests/integration/test_rag_existence_index.py @@ -34,7 +34,9 @@ TABLE = f"rag_{COLLECTION}" -async def _no_resolution(_name: str) -> None: +async def _no_resolution( + _name: str, _organization_id: object = None, _kb_id: object = None +) -> None: """A resolver that defers to the store's default embedder and width. `_ensure_collection` reads only the width to build the table; the embedder @@ -89,7 +91,7 @@ async def _clean_runtime_table(engine: AsyncEngine) -> AsyncGenerator[None, None async def test_ensure_collection_builds_an_index_per_lookup_key(engine: AsyncEngine) -> None: store = _store(engine) - await store._ensure_collection(COLLECTION) + await store._ensure_collection(COLLECTION, None) async with store.async_session() as session: result = await session.execute( @@ -118,7 +120,7 @@ async def test_source_path_wins_and_returns_that_documents_own_hash(engine: Asyn still win and hand back `live`'s own hash, not `decoy`'s. """ store = _store(engine) - await store._ensure_collection(COLLECTION) + await store._ensure_collection(COLLECTION, None) await _insert( store, doc_id="decoy", @@ -147,7 +149,7 @@ async def test_the_filename_fallback_keeps_the_unaddressed_rule(engine: AsyncEng """#990: a source_path miss matches a same-name document only where that document has not addressed itself under a different path.""" store = _store(engine) - await store._ensure_collection(COLLECTION) + await store._ensure_collection(COLLECTION, None) await _insert( store, doc_id="addressed", @@ -173,7 +175,7 @@ async def test_the_filename_fallback_keeps_the_unaddressed_rule(engine: AsyncEng async def test_content_hash_is_the_last_resort(engine: AsyncEngine) -> None: store = _store(engine) - await store._ensure_collection(COLLECTION) + await store._ensure_collection(COLLECTION, None) await _insert( store, doc_id="moved", source_path="/old/name.pdf", filename="name.pdf", content_hash="same" ) @@ -191,7 +193,7 @@ async def test_the_fallback_tiebreak_is_deterministic(engine: AsyncEngine) -> No chosen is fixed by `ORDER BY parent_doc_id, id`, not by heap order. Two unaddressed documents share a filename; the lower `parent_doc_id` wins.""" store = _store(engine) - await store._ensure_collection(COLLECTION) + await store._ensure_collection(COLLECTION, None) await _insert( store, doc_id="bbb", source_path="dup.pdf", filename="dup.pdf", content_hash="v-b" ) @@ -212,7 +214,7 @@ async def test_a_source_path_too_long_for_a_btree_index_still_ingests(engine: As long path would fail every ingest into the collection. The hash index has no such ceiling - this row inserts and is found.""" store = _store(engine) - await store._ensure_collection(COLLECTION) + await store._ensure_collection(COLLECTION, None) long_path = "s3://bucket/" + "a" * 3000 await _insert(store, doc_id="big", source_path=long_path, filename="big.pdf", content_hash="h") @@ -224,7 +226,7 @@ async def test_a_source_path_too_long_for_a_btree_index_still_ingests(engine: As async def test_no_key_matches_answers_none(engine: AsyncEngine) -> None: store = _store(engine) - await store._ensure_collection(COLLECTION) + await store._ensure_collection(COLLECTION, None) await _insert( store, doc_id="only", source_path="/srv/other.pdf", filename="other.pdf", content_hash="h" ) diff --git a/backend/tests/integration/test_vector_store_reserved_names.py b/backend/tests/integration/test_vector_store_reserved_names.py index 24e393dba..d449f637d 100644 --- a/backend/tests/integration/test_vector_store_reserved_names.py +++ b/backend/tests/integration/test_vector_store_reserved_names.py @@ -30,7 +30,9 @@ ) -async def _no_collection_of_its_own(name: str) -> None: +async def _no_collection_of_its_own( + name: str, organization_id: object = None, knowledge_base_id: object = None +) -> None: """A resolver that answers "nothing recorded for this one". Not `None` in place of the resolver: `resolver` is a required argument since @@ -90,7 +92,7 @@ async def test_a_collection_beside_it_is_created_and_dropped_for_real( """ store = _store_on(engine) - await store.create_collection("documents_archive") + await store.create_collection("documents_archive", organization_id=None) async with engine.connect() as connection: created = await connection.execute( diff --git a/backend/tests/test_capability_edges.py b/backend/tests/test_capability_edges.py index 820fae6c4..373bbd749 100644 --- a/backend/tests/test_capability_edges.py +++ b/backend/tests/test_capability_edges.py @@ -10,6 +10,7 @@ import asyncio from decimal import Decimal from unittest.mock import AsyncMock, MagicMock, patch +from uuid import uuid4 import pytest from pydantic_ai import ModelRetry, RunContext @@ -79,7 +80,7 @@ async def test_a_collection_that_fails_fails_the_whole_search(self): with pytest.raises(RuntimeError): await _retrieval_over(store).retrieve_multi( - query="anything", collection_names=["healthy", "broken"] + query="anything", collection_names=["healthy", "broken"], organization_id=None ) @pytest.mark.anyio @@ -91,7 +92,7 @@ async def test_an_empty_collection_is_not_a_failure(self): ) results = await _retrieval_over(store).retrieve_multi( - query="anything", collection_names=["populated", "never_ingested"] + query="anything", collection_names=["populated", "never_ingested"], organization_id=None ) assert [r.content for r in results] == ["found"] @@ -108,7 +109,7 @@ async def test_every_result_names_the_collection_it_came_from(self): store.search = AsyncMock(return_value=[SearchResult(content="chunk", score=0.5)]) results = await _retrieval_over(store).retrieve( - query="anything", collection_name="handbook" + query="anything", collection_name="handbook", organization_id=None ) assert [r.metadata["collection"] for r in results] == ["handbook"] @@ -118,7 +119,9 @@ class TestKnowledgeSearchGuards: @pytest.mark.anyio async def test_no_collections_says_so_rather_than_searching_everything(self): """The dangerous failure mode would be an unscoped search.""" - result = await search_knowledge_base(query="x", kb_collection_names=[]) + result = await search_knowledge_base( + query="x", kb_collection_names=[], organization_id=None + ) assert "No active knowledge bases" in result @pytest.mark.anyio @@ -129,7 +132,9 @@ async def test_one_collection_uses_the_single_collection_path(self): "app.agents.capabilities.knowledge._search.get_retrieval_service", return_value=service, ): - await search_knowledge_base(query="x", kb_collection_names=["kb_a"]) + await search_knowledge_base( + query="x", kb_collection_names=["kb_a"], organization_id=None + ) service.retrieve.assert_awaited_once() @pytest.mark.anyio @@ -140,7 +145,9 @@ async def test_several_collections_use_the_multi_path(self): "app.agents.capabilities.knowledge._search.get_retrieval_service", return_value=service, ): - await search_knowledge_base(query="x", kb_collection_names=["kb_a", "kb_b"]) + await search_knowledge_base( + query="x", kb_collection_names=["kb_a", "kb_b"], organization_id=None + ) service.retrieve_multi.assert_awaited_once() @pytest.mark.anyio @@ -155,7 +162,9 @@ async def test_a_retrieval_failure_surfaces_as_an_external_service_error(self): ), pytest.raises(ExternalServiceError), ): - await search_knowledge_base(query="x", kb_collection_names=["kb_a"]) + await search_knowledge_base( + query="x", kb_collection_names=["kb_a"], organization_id=None + ) @pytest.mark.anyio async def test_an_unconfigured_deployment_keeps_saying_what_to_configure(self): @@ -179,11 +188,70 @@ async def test_an_unconfigured_deployment_keeps_saying_what_to_configure(self): ), pytest.raises(ConfigurationError) as refusal, ): - await search_knowledge_base(query="x", kb_collection_names=["kb_a"]) + await search_knowledge_base( + query="x", kb_collection_names=["kb_a"], organization_id=None + ) assert refusal.value.details == {"setting": "OPENROUTER_API_KEY"} +class TestBoundKnowledgeBaseIds: + """The agent path carries the bound KB id, so a shared collection name + resolves the bound row's config and key, not another same-named row's (#913).""" + + @pytest.mark.anyio + async def test_a_single_search_pins_the_bound_kb_id(self): + service = MagicMock() + service.retrieve = AsyncMock(return_value=[]) + kb = uuid4() + with patch( + "app.agents.capabilities.knowledge._search.get_retrieval_service", + return_value=service, + ): + await search_knowledge_base( + query="x", + kb_collection_names=["kb_a"], + kb_collection_ids=[kb], + organization_id=None, + ) + assert service.retrieve.await_args.kwargs["knowledge_base_id"] == kb + + @pytest.mark.anyio + async def test_a_multi_search_pins_each_collections_bound_kb_id(self): + service = MagicMock() + service.retrieve_multi = AsyncMock(return_value=[]) + a, b = uuid4(), uuid4() + with patch( + "app.agents.capabilities.knowledge._search.get_retrieval_service", + return_value=service, + ): + await search_knowledge_base( + query="x", + kb_collection_names=["kb_a", "kb_b"], + kb_collection_ids=[a, b], + organization_id=None, + ) + assert service.retrieve_multi.await_args.kwargs["knowledge_base_ids"] == [a, b] + + @pytest.mark.anyio + async def test_ids_that_do_not_align_with_the_names_are_dropped(self): + """A length mismatch (the nameless ContextVar fallback, or a bug) resolves + by organization rather than pinning the wrong id to a collection.""" + service = MagicMock() + service.retrieve = AsyncMock(return_value=[]) + with patch( + "app.agents.capabilities.knowledge._search.get_retrieval_service", + return_value=service, + ): + await search_knowledge_base( + query="x", + kb_collection_names=["kb_a"], + kb_collection_ids=[uuid4(), uuid4()], + organization_id=None, + ) + assert service.retrieve.await_args.kwargs["knowledge_base_id"] is None + + class TestEmbeddingCredential: """An unconfigured embedding key is a configuration state, not a crash.""" @@ -255,7 +323,7 @@ def test_a_collection_nobody_has_uploaded_to_reports_as_empty(self): # empty collection must not depend on the query succeeding. store.async_session = MagicMock(side_effect=AssertionError("should not query")) - info = asyncio.run(store.get_collection_info("never_ingested")) + info = asyncio.run(store.get_collection_info("never_ingested", organization_id=None)) assert (info.name, info.total_vectors, info.dim) == ("never_ingested", 0, 1536) @@ -276,7 +344,7 @@ def test_searching_a_collection_nobody_has_uploaded_to_finds_nothing(self): store._for_collection = AsyncMock(side_effect=AssertionError("should not embed")) store.async_session = MagicMock(side_effect=AssertionError("should not query")) - assert asyncio.run(store.search("never_ingested", "anything")) == [] + assert asyncio.run(store.search("never_ingested", "anything", organization_id=None)) == [] def test_the_service_builds_on_a_deployment_with_no_key(self, monkeypatch): """`get_embedding_service` is a FastAPI dependency of every RAG route. diff --git a/backend/tests/test_chunk_insert_batching.py b/backend/tests/test_chunk_insert_batching.py index df675ffbc..a5f7e0c2a 100644 --- a/backend/tests/test_chunk_insert_batching.py +++ b/backend/tests/test_chunk_insert_batching.py @@ -98,7 +98,7 @@ async def test_a_document_of_many_chunks_is_written_in_batches(self, monkeypatch monkeypatch.setattr(vectorstore_module, "_CHUNK_INSERT_BATCH", 100) session = RecordingSession() - await _store(session).insert_document("docs", _document(chunks=250)) + await _store(session).insert_document("docs", _document(chunks=250), organization_id=None) assert len(session.calls) == 3 assert [len(batch) for batch in session.calls] == [100, 100, 50] # ty: ignore[invalid-argument-type] @@ -109,7 +109,7 @@ async def test_a_document_within_one_batch_is_one_statement(self, monkeypatch): monkeypatch.setattr(vectorstore_module, "_CHUNK_INSERT_BATCH", 200) session = RecordingSession() - await _store(session).insert_document("docs", _document(chunks=7)) + await _store(session).insert_document("docs", _document(chunks=7), organization_id=None) assert len(session.calls) == 1 assert len(session.calls[0]) == 7 # ty: ignore[invalid-argument-type] @@ -120,8 +120,8 @@ async def test_the_statement_count_does_not_grow_with_the_chunk_count(self, monk monkeypatch.setattr(vectorstore_module, "_CHUNK_INSERT_BATCH", 500) small, large = RecordingSession(), RecordingSession() - await _store(small).insert_document("docs", _document(chunks=40)) - await _store(large).insert_document("docs", _document(chunks=400)) + await _store(small).insert_document("docs", _document(chunks=40), organization_id=None) + await _store(large).insert_document("docs", _document(chunks=400), organization_id=None) assert len(small.calls) == len(large.calls) == 1 @@ -131,7 +131,7 @@ async def test_every_chunk_is_in_the_parameters_exactly_once(self, monkeypatch): session = RecordingSession() document = _document(chunks=10) - await _store(session).insert_document("docs", document) + await _store(session).insert_document("docs", document, organization_id=None) written = [row["id"] for batch in session.calls for row in batch] # ty: ignore[invalid-argument-type] assert written == [chunk.chunk_id for chunk in document.chunked_pages] @@ -145,7 +145,7 @@ async def test_each_row_carries_its_own_metadata_and_vector(self, monkeypatch): for index, chunk in enumerate(document.chunked_pages): chunk.page_num = index + 1 - await _store(session).insert_document("docs", document) + await _store(session).insert_document("docs", document, organization_id=None) rows = session.calls[0] assert [json.loads(row["metadata"])["page_num"] for row in rows] == [1, 2, 3, 4] # ty: ignore[invalid-argument-type] @@ -178,7 +178,7 @@ def counting(chunk: object, document: object) -> dict[str, object]: seen: list[int] = [] session.on_execute = lambda: seen.append(built) - await store.insert_document("docs", _document(chunks=30)) + await store.insert_document("docs", _document(chunks=30), organization_id=None) assert seen == [10, 20, 30] @@ -188,6 +188,6 @@ async def test_a_document_with_no_chunks_is_refused_before_any_statement(self): session = RecordingSession() with pytest.raises(ValueError, match="no chunked pages"): - await _store(session).insert_document("docs", _document(chunks=0)) + await _store(session).insert_document("docs", _document(chunks=0), organization_id=None) assert session.calls == [] diff --git a/backend/tests/test_collection_name_rules.py b/backend/tests/test_collection_name_rules.py index 921eb6745..0e9d502ae 100644 --- a/backend/tests/test_collection_name_rules.py +++ b/backend/tests/test_collection_name_rules.py @@ -212,6 +212,6 @@ async def test_a_name_creation_refuses_cannot_reach_the_drop_path_either( store = PgVectorStore.__new__(PgVectorStore) with pytest.raises(BadRequestError): - await store.create_collection(name) + await store.create_collection(name, organization_id=None) with pytest.raises(BadRequestError): await store.delete_collection(name) diff --git a/backend/tests/test_commands.py b/backend/tests/test_commands.py index 91d7812aa..0540a3bbe 100644 --- a/backend/tests/test_commands.py +++ b/backend/tests/test_commands.py @@ -447,6 +447,28 @@ async def session() -> AsyncGenerator[object, None]: assert db.begin_nested.call_count == 2 +class TestRagSearchCommand: + """`rag-search` is a tenantless operator path. + + `RetrievalService.retrieve` takes `organization_id` as a required + keyword-only argument to scope reranker resolution. The CLI has no acting + tenant, so it must pass `organization_id=None` explicitly; omitting it + raised `TypeError` before any search ran. + """ + + def test_search_async_scopes_retrieval_to_no_organization(self): + from unittest.mock import AsyncMock + + from app.commands import rag as rag_command + + retrieval = AsyncMock() + retrieval.retrieve = AsyncMock(return_value=[]) + + asyncio.run(rag_command.search_async("q", "handbook", 4, retrieval)) + + assert retrieval.retrieve.await_args.kwargs["organization_id"] is None + + class TestTheConsoleScript: """`agenticos` lives in `cli/`, which the unit suite never imported, so its dependencies were only ever exercised by the e2e seed step - 45 seconds into diff --git a/backend/tests/test_coverage_gate.py b/backend/tests/test_coverage_gate.py index 9debb1f97..f1a133765 100644 --- a/backend/tests/test_coverage_gate.py +++ b/backend/tests/test_coverage_gate.py @@ -86,8 +86,10 @@ def _matches_glob(path: str, pattern: str) -> bool: "app/services/collection_access.py", "app/services/embed_session.py", "app/services/embedding_resolution.py", + "app/services/rerank_resolution.py", "app/services/health.py", "app/services/ingestion_config.py", + "app/services/knowledge_search.py", "app/services/mcp_catalog.py", "app/services/mcp_connection.py", "app/services/model_profile.py", diff --git a/backend/tests/test_embedding_resolution.py b/backend/tests/test_embedding_resolution.py index 7853463ac..b5dde8294 100644 --- a/backend/tests/test_embedding_resolution.py +++ b/backend/tests/test_embedding_resolution.py @@ -80,9 +80,9 @@ async def _resolve(kb, secret_row=None): ): db_ctx.return_value.__aenter__ = AsyncMock(return_value=MagicMock()) db_ctx.return_value.__aexit__ = AsyncMock(return_value=False) - kbs.get_by_collection_name = AsyncMock(return_value=kb) + kbs.get_for_collection = AsyncMock(return_value=kb) secrets.get = AsyncMock(return_value=secret_row) - return await embeddings_for_collection("handbook"), secrets + return await embeddings_for_collection("handbook", organization_id=uuid.uuid4()), secrets class TestResolution: @@ -129,6 +129,34 @@ async def test_a_repr_never_carries_the_key(self): assert "organization" in repr(resolved) +class TestAuthorizedKnowledgeBaseId: + """A given `knowledge_base_id` resolves that exact row, not a name lookup, so + a shared collection name cannot re-select and unseal another row's key (#913).""" + + async def test_a_given_kb_id_reads_that_row_and_skips_the_name_lookup(self): + with ( + patch(f"{_MODULE}.get_db_context") as db_ctx, + patch(f"{_MODULE}.knowledge_base_repo") as kbs, + patch(f"{_MODULE}.organization_secret_repo") as secrets, + patch(f"{_MODULE}.settings") as env, + ): + env.OPENROUTER_API_KEY = "sk-deployment" + db_ctx.return_value.__aenter__ = AsyncMock(return_value=MagicMock()) + db_ctx.return_value.__aexit__ = AsyncMock(return_value=False) + kbs.get_by_id = AsyncMock(return_value=_kb()) + kbs.get_for_collection = AsyncMock(return_value=_kb(secret_id=uuid.uuid4())) + secrets.get = AsyncMock() + kb_id = uuid.uuid4() + + resolved = await embeddings_for_collection( + "handbook", organization_id=_ORG, knowledge_base_id=kb_id + ) + + kbs.get_by_id.assert_awaited_once_with(kbs.get_by_id.await_args.args[0], kb_id) + kbs.get_for_collection.assert_not_called() + assert resolved is not None and resolved.key_source == EmbeddingKeySource.DEPLOYMENT + + class TestCredentialDegradation: """Every failure lands on the deployment key, saying which failure it was. diff --git a/backend/tests/test_ingestion_embedding_key.py b/backend/tests/test_ingestion_embedding_key.py index caa21e723..2d2d816f9 100644 --- a/backend/tests/test_ingestion_embedding_key.py +++ b/backend/tests/test_ingestion_embedding_key.py @@ -106,7 +106,7 @@ async def _store() -> PgVectorStore: "app.worker.tasks.rag_tasks.create_async_engine", return_value=MagicMock(dispose=AsyncMock()), ): - async with _ingestion_service(processor=MagicMock()) as service: + async with _ingestion_service(processor=MagicMock(), organization_id=None) as service: store = service.store assert isinstance(store, PgVectorStore) return store @@ -163,12 +163,12 @@ async def _the_flows_embedder( db_ctx.return_value.__aenter__ = AsyncMock(return_value=MagicMock()) db_ctx.return_value.__aexit__ = AsyncMock(return_value=False) - bases.get_by_collection_name = AsyncMock(return_value=_knowledge_base(secret_id=secret_id)) + bases.get_for_collection = AsyncMock(return_value=_knowledge_base(secret_id=secret_id)) secrets.get = AsyncMock(return_value=vault_row) resolution_env.OPENROUTER_API_KEY = deployment_key embedding_env.OPENROUTER_API_KEY = deployment_key - embedder, dim = await (await _store())._for_collection("handbook") + embedder, dim = await (await _store())._for_collection("handbook", None) yield embedder, dim, openai @@ -239,11 +239,13 @@ async def test_two_collections_on_one_key_do_not_share_each_others_name(self): "handbook": _resolved(EmbeddingKeySource.SECRET_MISSING), "policies": _resolved(EmbeddingKeySource.DEPLOYMENT), } - store._resolver = AsyncMock(side_effect=lambda name: resolutions[name]) + store._resolver = AsyncMock( + side_effect=lambda name, organization_id=None, knowledge_base_id=None: resolutions[name] + ) origins = [] for collection in resolutions: - embedder, _ = await store._for_collection(collection) + embedder, _ = await store._for_collection(collection, None) with pytest.raises(ConfigurationError) as refusal: embedder.embed_query("anything") origins.append(refusal.value.details["key_origin"]) @@ -332,7 +334,7 @@ async def _resolution(self, key_source: EmbeddingKeySource): ), patch("app.worker.tasks.rag_tasks._say_in_flow_log") as said, ): - answer = await _announcing_resolver()("handbook") + answer = await _announcing_resolver()("handbook", None) return answer, said async def test_a_degraded_credential_is_announced_with_the_reason(self): @@ -361,7 +363,7 @@ async def test_a_collection_no_knowledge_base_claims_says_nothing(self): ), patch("app.worker.tasks.rag_tasks._say_in_flow_log") as said, ): - assert await _announcing_resolver()("unclaimed") is None + assert await _announcing_resolver()("unclaimed", None) is None said.assert_not_called() @@ -376,13 +378,17 @@ async def test_one_collection_is_announced_once_however_often_it_resolves(self): with ( patch( "app.worker.tasks.rag_tasks.embeddings_for_collection", - new=AsyncMock(side_effect=lambda name: resolutions[name]), + new=AsyncMock( + side_effect=lambda name, organization_id=None, knowledge_base_id=None: ( + resolutions[name] + ) + ), ), patch("app.worker.tasks.rag_tasks._say_in_flow_log") as said, ): resolve = _announcing_resolver() for name in ("handbook", "handbook", "policies", "handbook"): - await resolve(name) + await resolve(name, None) assert [call.args[0].split("'")[1] for call in said.call_args_list] == [ "handbook", @@ -401,8 +407,8 @@ async def test_a_second_flow_run_reports_a_credential_that_is_still_broken(self) ), patch("app.worker.tasks.rag_tasks._say_in_flow_log") as said, ): - await _announcing_resolver()("handbook") - await _announcing_resolver()("handbook") + await _announcing_resolver()("handbook", None) + await _announcing_resolver()("handbook", None) assert said.call_count == 2 diff --git a/backend/tests/test_ingestion_spend_repo.py b/backend/tests/test_ingestion_spend_repo.py index 3f0b25ebb..50441c544 100644 --- a/backend/tests/test_ingestion_spend_repo.py +++ b/backend/tests/test_ingestion_spend_repo.py @@ -15,7 +15,7 @@ import pytest from sqlalchemy.dialects import postgresql -from app.db.models.ingestion_spend import IngestionSpend +from app.db.models.ingestion_spend import IngestionSpend, SpendSource from app.repositories import ingestion_spend_repo pytestmark = pytest.mark.anyio @@ -84,6 +84,35 @@ async def test_an_unpriced_window_is_flagged_not_silently_free(self): assert spend.cost_is_partial is True + async def test_a_row_is_indexing_unless_told_otherwise(self): + session = _RecordingSession() + spend = await ingestion_spend_repo.record( + session, + organization_id=None, + rag_document_id=None, + model="text-embedding-3-large", + input_tokens=1, + output_tokens=0, + cost_usd=Decimal(0), + cost_is_partial=False, + ) + assert spend.source == SpendSource.INGESTION.value + + async def test_a_search_is_recorded_as_retrieval(self): + session = _RecordingSession() + spend = await ingestion_spend_repo.record( + session, + organization_id=None, + rag_document_id=None, + model="rerank-v3.5", + input_tokens=0, + output_tokens=0, + cost_usd=Decimal("0.002"), + cost_is_partial=False, + source=SpendSource.RETRIEVAL, + ) + assert spend.source == SpendSource.RETRIEVAL.value + class TestSumming: async def test_the_sum_filters_on_the_organization_and_the_window(self): @@ -135,3 +164,23 @@ async def test_a_window_with_no_ingestion_sums_to_zero_not_none(self): ) assert total == Decimal(0) + + async def test_a_source_narrows_the_window_and_none_does_not(self): + # The dashboard sums each source on its own so a search is not counted + # as indexing; the budget sums both, so an unfiltered call must not carry + # the predicate. + session = _RecordingSession(scalar_result=Decimal("0.30")) + window = { + "organization_id": uuid.uuid4(), + "start": datetime(2026, 7, 1, tzinfo=UTC), + "end": datetime(2026, 8, 1, tzinfo=UTC), + } + + await ingestion_spend_repo.sum_cost_window(session, **window, source=SpendSource.RETRIEVAL) + narrowed = session.statements[-1].compile(dialect=postgresql.dialect()).params + assert SpendSource.RETRIEVAL.value in narrowed.values() + + await ingestion_spend_repo.sum_cost_window(session, **window) + unfiltered = session.statements[-1].compile(dialect=postgresql.dialect()).params + assert SpendSource.RETRIEVAL.value not in unfiltered.values() + assert SpendSource.INGESTION.value not in unfiltered.values() diff --git a/backend/tests/test_kb_scoping.py b/backend/tests/test_kb_scoping.py index 5bd5f26df..ec2f81dae 100644 --- a/backend/tests/test_kb_scoping.py +++ b/backend/tests/test_kb_scoping.py @@ -62,6 +62,10 @@ def _kb( kb.organization_id = organization_id kb.is_default = is_default kb.visibility = visibility + kb.embedding_provider = "openrouter" + kb.embedding_model = "text-embedding-3-small" + kb.embedding_dim = 1536 + kb.embedding_secret_id = None return kb @@ -840,3 +844,175 @@ async def test_the_listing_gives_each_row_its_own_counts(self, mock_db): assert listing.total == 2 assert [item.document_count for item in listing.items] == [0, 3] assert [item.chunk_count for item in listing.items] == [0, 90] + + +class TestRerankConfig: + """Setting a collection's reranker, and the ways it is refused. + + Reranking is a model and a key together; a lone half reads as configured + and does nothing, and a key of the wrong purpose bills nobody's reranking. + Both are refused where the person setting them can see why.""" + + @pytest.fixture + def mock_db(self): + return MagicMock() + + def _cohere_secret(self, purpose: str = "cohere") -> MagicMock: + return MagicMock(purpose=purpose) + + @pytest.mark.anyio + async def test_a_model_without_a_key_is_refused(self, mock_db, unclaimed_collection_name): + data = KnowledgeBaseCreate( + name="KB", scope="org", collection_name="c", rerank_model="rerank-v3.5" + ) + with pytest.raises(BadRequestError, match="both a model and a key"): + await KnowledgeBaseService(mock_db).create(data, ctx=_ctx()) + + @pytest.mark.anyio + async def test_a_key_without_a_model_is_refused(self, mock_db, unclaimed_collection_name): + data = KnowledgeBaseCreate( + name="KB", scope="org", collection_name="c", rerank_secret_id=uuid.uuid4() + ) + with pytest.raises(BadRequestError, match="both a model and a key"): + await KnowledgeBaseService(mock_db).create(data, ctx=_ctx()) + + @pytest.mark.anyio + async def test_an_unsupported_model_is_refused_before_the_key_is_read( + self, mock_db, unclaimed_collection_name + ): + # A typo'd model with an otherwise valid key would be stored and shown as + # configured, then fail every search inside Cohere where the error is + # swallowed - reranking silently off. Refused at create, and before the + # vault is even consulted. + data = KnowledgeBaseCreate( + name="KB", + scope="org", + collection_name="c", + rerank_model="rerank-v3.5x", + rerank_secret_id=uuid.uuid4(), + ) + with ( + patch("app.repositories.organization_secret_repo.get", new=AsyncMock()) as secret_get, + pytest.raises(BadRequestError, match="Unsupported rerank model"), + ): + await KnowledgeBaseService(mock_db).create(data, ctx=_ctx()) + secret_get.assert_not_called() + + @pytest.mark.anyio + async def test_an_update_to_an_unsupported_model_is_refused(self, mock_db): + kb = _kb("org", organization_id=uuid.uuid4()) + data = KnowledgeBaseUpdate(rerank_model="bogus", rerank_secret_id=uuid.uuid4()) + with ( + patch.object(KnowledgeBaseService, "get_for_write", new=AsyncMock(return_value=kb)), + pytest.raises(BadRequestError, match="Unsupported rerank model"), + ): + await KnowledgeBaseService(mock_db).update(kb.id, data, ctx=_ctx()) + + @pytest.mark.anyio + async def test_a_configured_pair_is_written_through(self, mock_db, unclaimed_collection_name): + secret_id = uuid.uuid4() + data = KnowledgeBaseCreate( + name="KB", + scope="org", + collection_name="c", + rerank_model="rerank-v3.5", + rerank_secret_id=secret_id, + ) + with ( + patch( + "app.repositories.organization_secret_repo.get", + new=AsyncMock(return_value=self._cohere_secret()), + ), + patch( + "app.services.knowledge_base.resolve_access", + new=AsyncMock(return_value=True), + ), + patch( + "app.repositories.knowledge_base_repo.create", + new=AsyncMock(return_value=MagicMock()), + ) as created, + ): + await KnowledgeBaseService(mock_db).create(data, ctx=_ctx()) + + assert created.call_args.kwargs["rerank_model"] == "rerank-v3.5" + assert created.call_args.kwargs["rerank_secret_id"] == secret_id + + @pytest.mark.anyio + async def test_a_key_the_caller_cannot_reach_is_refused_as_missing( + self, mock_db, unclaimed_collection_name + ): + # In the organization's vault, but private to another member: binding it + # would lend a key `secrets:view` refuses the caller. Refused as a miss, + # so the refusal cannot be told from "no such key" and used to enumerate. + data = KnowledgeBaseCreate( + name="KB", + scope="org", + collection_name="c", + rerank_model="rerank-v3.5", + rerank_secret_id=uuid.uuid4(), + ) + with ( + patch( + "app.repositories.organization_secret_repo.get", + new=AsyncMock(return_value=self._cohere_secret()), + ), + patch( + "app.services.knowledge_base.resolve_access", + new=AsyncMock(return_value=False), + ), + pytest.raises(BadRequestError, match="not in this organization's vault"), + ): + await KnowledgeBaseService(mock_db).create(data, ctx=_ctx()) + + @pytest.mark.anyio + async def test_a_key_of_the_wrong_purpose_is_refused(self, mock_db, unclaimed_collection_name): + data = KnowledgeBaseCreate( + name="KB", + scope="org", + collection_name="c", + rerank_model="rerank-v3.5", + rerank_secret_id=uuid.uuid4(), + ) + with ( + patch( + "app.repositories.organization_secret_repo.get", + new=AsyncMock(return_value=self._cohere_secret(purpose="openrouter")), + ), + patch( + "app.services.knowledge_base.resolve_access", + new=AsyncMock(return_value=True), + ), + pytest.raises(BadRequestError, match="reranking runs through"), + ): + await KnowledgeBaseService(mock_db).create(data, ctx=_ctx()) + + @pytest.mark.anyio + async def test_an_update_turns_reranking_off_by_sending_both_null(self, mock_db): + kb = _kb("org", organization_id=uuid.uuid4()) + data = KnowledgeBaseUpdate(rerank_model=None, rerank_secret_id=None) + with ( + patch.object(KnowledgeBaseService, "get_for_write", new=AsyncMock(return_value=kb)), + patch( + "app.repositories.knowledge_base_repo.update", + new=AsyncMock(return_value=kb), + ) as updated, + ): + await KnowledgeBaseService(mock_db).update(kb.id, data, ctx=_ctx()) + + assert updated.call_args.kwargs["set_rerank"] is True + assert updated.call_args.kwargs["rerank_model"] is None + + @pytest.mark.anyio + async def test_an_update_about_something_else_leaves_reranking_alone(self, mock_db): + kb = _kb("org", organization_id=uuid.uuid4()) + data = KnowledgeBaseUpdate(name="Renamed") + with ( + patch.object(KnowledgeBaseService, "get_for_write", new=AsyncMock(return_value=kb)), + patch( + "app.repositories.knowledge_base_repo.update", + new=AsyncMock(return_value=kb), + ) as updated, + ): + await KnowledgeBaseService(mock_db).update(kb.id, data, ctx=_ctx()) + + assert updated.call_args.kwargs["set_rerank"] is False diff --git a/backend/tests/test_knowledge_base_repo.py b/backend/tests/test_knowledge_base_repo.py new file mode 100644 index 000000000..31501416c --- /dev/null +++ b/backend/tests/test_knowledge_base_repo.py @@ -0,0 +1,67 @@ +"""Resolving a collection name to the right organization's knowledge base. + +`collection_name` is not unique across tenants, so resolving one by name alone +can return another organization's row - and then unseal and bill that +organization's key (#913). `get_for_collection` narrows by organization in two +passes: the caller's own row wins, and an `app`-scoped row (owned by no +organization) is the shared fallback. +""" + +import uuid +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from app.repositories import knowledge_base as kb_repo + +pytestmark = pytest.mark.anyio + + +def _kb(organization_id: uuid.UUID | None) -> MagicMock: + kb = MagicMock() + kb.organization_id = organization_id + return kb + + +async def _resolve(candidates: list[MagicMock], organization_id: uuid.UUID | None) -> MagicMock: + with patch.object(kb_repo, "list_by_collection_name", new=AsyncMock(return_value=candidates)): + return await kb_repo.get_for_collection(MagicMock(), "shared", organization_id) + + +async def test_a_shared_name_resolves_to_the_callers_own_organization() -> None: + org_a, org_b = uuid.uuid4(), uuid.uuid4() + a, b = _kb(org_a), _kb(org_b) + + assert await _resolve([a, b], org_a) is a + assert await _resolve([a, b], org_b) is b + + +async def test_an_organization_without_its_own_row_never_gets_anothers() -> None: + """The security property: org B resolving a name only org A holds gets + nothing, not org A's key.""" + only_a = _kb(uuid.uuid4()) + + assert await _resolve([only_a], uuid.uuid4()) is None + + +async def test_an_app_scoped_row_is_the_shared_fallback() -> None: + """A collection owned by no organization is matched on the second pass, so an + organization with no row of its own still resolves the deployment-wide one.""" + app_kb = _kb(None) + + assert await _resolve([app_kb], uuid.uuid4()) is app_kb + + +async def test_an_organizations_own_row_wins_over_an_app_scoped_one() -> None: + org = uuid.uuid4() + app_kb, own = _kb(None), _kb(org) + + assert await _resolve([app_kb, own], org) is own + + +async def test_no_tenant_takes_the_first_candidate() -> None: + """`organization_id=None` is a CLI ingest with no tenant to scope to, where + the old name-only behaviour stands.""" + first, second = _kb(uuid.uuid4()), _kb(uuid.uuid4()) + + assert await _resolve([first, second], None) is first diff --git a/backend/tests/test_knowledge_search.py b/backend/tests/test_knowledge_search.py new file mode 100644 index 000000000..5b6f198ca --- /dev/null +++ b/backend/tests/test_knowledge_search.py @@ -0,0 +1,211 @@ +"""The knowledge-search request path, and that it meters what it spends. + +The search itself is delegated to retrieval; what this service adds is a ledger +scoped to the organization, opened around the search so the embedding and any +rerank book to it, and persisted afterwards. So the tests drive a retrieval +double that books spend the way an embedding or a rerank would, and assert it +reaches `ingestion_spend`. +""" + +from __future__ import annotations + +import uuid +from decimal import Decimal +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from app.agents.capabilities.budget import ( + BudgetExceeded, + BudgetScope, + SpendEntry, + book_ambient_spend, +) +from app.db.models.ingestion_spend import SpendSource +from app.schemas.rag import RAGSearchRequest +from app.services.knowledge_search import KnowledgeSearchService +from app.services.rag.models import SearchResult + +pytestmark = pytest.mark.anyio + +_MODULE = "app.services.knowledge_search" + + +@pytest.fixture(autouse=True) +def _budget_ok(): + """Every test but the budget one runs under an organization within its cap.""" + with patch(f"{_MODULE}.assert_organization_within_budget", new=AsyncMock()): + yield + + +def _kb(collection_name: str) -> MagicMock: + return MagicMock(collection_name=collection_name) + + +def _ctx(organization_id: uuid.UUID | None = None) -> MagicMock: + return MagicMock(organization_id=organization_id or uuid.uuid4()) + + +def _service(*, readable: list[MagicMock], retrieve=None, retrieve_multi=None): + access = MagicMock() + access.readable_all = AsyncMock(return_value=readable) + retrieval = MagicMock() + retrieval.retrieve = AsyncMock(side_effect=retrieve or _no_spend_retrieve) + retrieval.retrieve_multi = AsyncMock(side_effect=retrieve_multi or _no_spend_retrieve) + return KnowledgeSearchService(MagicMock(), retrieval, access), retrieval + + +async def _no_spend_retrieve(*args, **kwargs) -> list[SearchResult]: + return [SearchResult(content="hit", score=0.9)] + + +def _booking(cost: str, *, priced: bool = True, model: str = "rerank-v3.5"): + async def _retrieve(*args, **kwargs) -> list[SearchResult]: + book_ambient_spend( + SpendEntry( + model_name=model, + input_tokens=0, + output_tokens=0, + cost_usd=Decimal(cost), + priced=priced, + ) + ) + return [SearchResult(content="hit", score=0.9)] + + return _retrieve + + +class TestRouting: + async def test_one_collection_uses_the_single_collection_path(self): + service, retrieval = _service(readable=[_kb("c1")]) + with patch(f"{_MODULE}.ingestion_spend_repo.record", new=AsyncMock()): + await service.search(_ctx(), RAGSearchRequest(query="q", collection_name="c1")) + retrieval.retrieve.assert_awaited_once() + retrieval.retrieve_multi.assert_not_awaited() + + async def test_several_collections_use_the_multi_path(self): + service, retrieval = _service(readable=[_kb("a"), _kb("b")]) + with patch(f"{_MODULE}.ingestion_spend_repo.record", new=AsyncMock()): + await service.search(_ctx(), RAGSearchRequest(query="q", collection_names=["a", "b"])) + retrieval.retrieve_multi.assert_awaited_once() + retrieval.retrieve.assert_not_awaited() + + async def test_it_returns_what_retrieval_found(self): + service, _ = _service(readable=[_kb("c1")]) + with patch(f"{_MODULE}.ingestion_spend_repo.record", new=AsyncMock()): + results = await service.search(_ctx(), RAGSearchRequest(query="q")) + assert [r.content for r in results] == ["hit"] + + +class TestMetering: + async def test_spend_booked_during_the_search_is_persisted_to_the_organization(self): + org = uuid.uuid4() + service, _ = _service(readable=[_kb("c1")], retrieve=_booking("0.002")) + with patch(f"{_MODULE}.ingestion_spend_repo.record", new=AsyncMock()) as record: + await service.search(_ctx(org), RAGSearchRequest(query="q")) + + record.assert_awaited_once() + kwargs = record.await_args.kwargs + assert kwargs["organization_id"] == org + assert kwargs["rag_document_id"] is None + assert kwargs["cost_usd"] == Decimal("0.002") + assert kwargs["cost_is_partial"] is False + assert kwargs["source"] is SpendSource.RETRIEVAL + + async def test_a_search_that_spends_nothing_writes_no_row(self): + service, _ = _service(readable=[_kb("c1")]) + with patch(f"{_MODULE}.ingestion_spend_repo.record", new=AsyncMock()) as record: + await service.search(_ctx(), RAGSearchRequest(query="q")) + record.assert_not_awaited() + + async def test_an_unpriced_entry_makes_the_recorded_cost_partial(self): + service, _ = _service( + readable=[_kb("c1")], retrieve=_booking("0", priced=False, model="mystery") + ) + with patch(f"{_MODULE}.ingestion_spend_repo.record", new=AsyncMock()) as record: + await service.search(_ctx(), RAGSearchRequest(query="q")) + assert record.await_args.kwargs["cost_is_partial"] is True + + async def test_spend_in_two_models_is_one_row_each(self): + async def _retrieve(*args, **kwargs) -> list[SearchResult]: + book_ambient_spend(SpendEntry("text-embedding-3-large", 100, 0, Decimal("0.001"), True)) + book_ambient_spend(SpendEntry("rerank-v3.5", 0, 0, Decimal("0.002"), True)) + return [] + + service, _ = _service(readable=[_kb("c1")], retrieve=_retrieve) + with patch(f"{_MODULE}.ingestion_spend_repo.record", new=AsyncMock()) as record: + await service.search(_ctx(), RAGSearchRequest(query="q")) + + assert record.await_count == 2 + models = {call.kwargs["model"] for call in record.await_args_list} + assert models == {"text-embedding-3-large", "rerank-v3.5"} + + +class TestBudgetGuardsThePaidSearch: + """A search is a paid call, so the monthly cap refuses it before it spends, + exactly as ingestion is refused - not merely recorded after the fact.""" + + async def test_an_exhausted_budget_refuses_before_any_paid_call(self): + service, retrieval = _service(readable=[_kb("c1")], retrieve=_booking("0.002")) + exceeded = AsyncMock( + side_effect=BudgetExceeded( + limit_usd=Decimal("1"), spent_usd=Decimal("1"), scope=BudgetScope.ORGANIZATION + ) + ) + with ( + patch(f"{_MODULE}.assert_organization_within_budget", new=exceeded), + patch(f"{_MODULE}.ingestion_spend_repo.record", new=AsyncMock()) as record, + pytest.raises(BudgetExceeded), + ): + await service.search(_ctx(), RAGSearchRequest(query="q")) + + retrieval.retrieve.assert_not_awaited() + record.assert_not_awaited() + + async def test_a_search_with_no_organization_is_not_budget_checked(self): + service, _ = _service(readable=[_kb("c1")]) + ctx = MagicMock(organization_id=None) + with patch(f"{_MODULE}.assert_organization_within_budget") as guard: + await service.search(ctx, RAGSearchRequest(query="q")) + guard.assert_not_called() + + +class TestSpendSurvivesAFailedSearch: + """The query embedding is booked before the vector query it pays for, so a + search that fails mid-flight has already spent. That spend is booked through + a session of its own, because the request's own is about to roll back.""" + + async def test_spend_before_a_failure_is_persisted_out_of_band(self): + org = uuid.uuid4() + + async def _book_then_fail(*args, **kwargs) -> list[SearchResult]: + book_ambient_spend(SpendEntry("text-embedding-3-large", 100, 0, Decimal("0.001"), True)) + raise RuntimeError("vector store down") + + service, _ = _service(readable=[_kb("c1")], retrieve=_book_then_fail) + fresh_db = MagicMock() + fresh_session = MagicMock() + fresh_session.__aenter__ = AsyncMock(return_value=fresh_db) + fresh_session.__aexit__ = AsyncMock(return_value=False) + with ( + patch(f"{_MODULE}.get_db_context", return_value=fresh_session), + patch(f"{_MODULE}.ingestion_spend_repo.record", new=AsyncMock()) as record, + pytest.raises(RuntimeError, match="vector store down"), + ): + await service.search(_ctx(org), RAGSearchRequest(query="q")) + + record.assert_awaited_once() + assert record.await_args.args[0] is fresh_db + assert record.await_args.kwargs["organization_id"] == org + + async def test_a_failure_that_spent_nothing_opens_no_session(self): + async def _fail(*args, **kwargs) -> list[SearchResult]: + raise RuntimeError("down before the embedding") + + service, _ = _service(readable=[_kb("c1")], retrieve=_fail) + with ( + patch(f"{_MODULE}.get_db_context") as db_ctx, + pytest.raises(RuntimeError), + ): + await service.search(_ctx(), RAGSearchRequest(query="q")) + db_ctx.assert_not_called() diff --git a/backend/tests/test_rag_chunk_count.py b/backend/tests/test_rag_chunk_count.py index d91c8b04a..d4a54678d 100644 --- a/backend/tests/test_rag_chunk_count.py +++ b/backend/tests/test_rag_chunk_count.py @@ -58,7 +58,7 @@ def _document(*, chunks: int) -> Document: def _service(processor: MagicMock) -> IngestionService: store = MagicMock(insert_document=AsyncMock(), delete_document=AsyncMock()) store.find_existing_document = BaseVectorStore.find_existing_document.__get__(store) - return IngestionService(processor=processor, vector_store=store) + return IngestionService(processor=processor, vector_store=store, organization_id=None) class TestWhatThePipelineReports: diff --git a/backend/tests/test_rag_document_lookup.py b/backend/tests/test_rag_document_lookup.py index c19477350..7df72d7b3 100644 --- a/backend/tests/test_rag_document_lookup.py +++ b/backend/tests/test_rag_document_lookup.py @@ -57,7 +57,7 @@ def _store(docs: list[DocumentInfo], **overrides: object) -> MagicMock: def _service(docs: list[DocumentInfo]) -> IngestionService: - return IngestionService(processor=MagicMock(), vector_store=_store(docs)) + return IngestionService(processor=MagicMock(), vector_store=_store(docs), organization_id=None) def _doc( @@ -146,7 +146,7 @@ async def test_a_store_failure_answers_no_match(self): """A listing that cannot be read is not evidence the document is absent - but treating it as a match would delete one on a failed query.""" store = _store([], get_documents=AsyncMock(side_effect=RuntimeError("connection refused"))) - service = IngestionService(processor=MagicMock(), vector_store=store) + service = IngestionService(processor=MagicMock(), vector_store=store, organization_id=None) assert await service.existing_document("kb", "/srv/sync/handbook.pdf") == StoredDocument() @@ -232,7 +232,7 @@ async def test_both_answers_cost_one_read(self): ) ] ) - service = IngestionService(processor=MagicMock(), vector_store=store) + service = IngestionService(processor=MagicMock(), vector_store=store, organization_id=None) existing = await service.existing_document("kb", "/srv/sync/handbook.pdf") @@ -255,7 +255,7 @@ async def test_an_ingest_that_replaces_reads_the_collection_once(self): ] document.metadata.content_hash = "hash-a" processor = MagicMock(process_file=AsyncMock(return_value=document)) - service = IngestionService(processor=processor, vector_store=store) + service = IngestionService(processor=processor, vector_store=store, organization_id=None) result = await service.ingest_file( filepath=Path("handbook.pdf"), @@ -301,7 +301,9 @@ def _replacing(insert: AsyncMock) -> tuple[MagicMock, IngestionService]: ] document.metadata.content_hash = "hash-new" processor = MagicMock(process_file=AsyncMock(return_value=document)) - return store, IngestionService(processor=processor, vector_store=store) + return store, IngestionService( + processor=processor, vector_store=store, organization_id=None + ) async def test_a_failed_embedding_leaves_the_old_document_in_place(self): store, service = self._replacing(AsyncMock(side_effect=RuntimeError("provider refused"))) diff --git a/backend/tests/test_rag_failure_messages.py b/backend/tests/test_rag_failure_messages.py index f8a8de966..734cad178 100644 --- a/backend/tests/test_rag_failure_messages.py +++ b/backend/tests/test_rag_failure_messages.py @@ -149,7 +149,7 @@ class TestTheIngestionServiceMovesItToTheLog: def _service(*, parse: AsyncMock, insert: AsyncMock) -> IngestionService: processor = MagicMock(process_file=parse) store = MagicMock(insert_document=insert) - return IngestionService(processor=processor, vector_store=store) + return IngestionService(processor=processor, vector_store=store, organization_id=None) @staticmethod def _document() -> MagicMock: diff --git a/backend/tests/test_rerank_resolution.py b/backend/tests/test_rerank_resolution.py new file mode 100644 index 000000000..44bd6b7ea --- /dev/null +++ b/backend/tests/test_rerank_resolution.py @@ -0,0 +1,200 @@ +"""Tests for per-collection reranker resolution. + +The mirror image of embedding resolution, and its one deliberate difference +carries the weight here: where a missing embedding key falls back to the +deployment's, a missing rerank key turns reranking *off*. So every path but a +usable organization secret resolves to `None`, and the three that mean a real +misconfiguration say so in a log line while the normal off state stays silent. +""" + +from __future__ import annotations + +import logging +import uuid +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from pydantic import SecretStr + +from app.core.secret_kinds import ApiKeySecret, SecretKind, seal_secret +from app.core.vault import VaultScope +from app.services.rerank_resolution import ( + RerankKeySource, + ResolvedReranker, + reranker_for_collection, +) + +pytestmark = pytest.mark.anyio + +_MODULE = "app.services.rerank_resolution" +_ORG = uuid.uuid4() + + +def _kb( + *, + model: str | None = "rerank-v3.5", + secret_id: uuid.UUID | None = None, + organization_id: uuid.UUID | None = _ORG, +): + return MagicMock( + collection_name="handbook", + rerank_model=model, + rerank_secret_id=secret_id, + organization_id=organization_id, + ) + + +def _sealed_key_row(plaintext: str, *, organization_id: uuid.UUID = _ORG): + sealed = seal_secret( + ApiKeySecret(api_key=SecretStr(plaintext)), + scope=VaultScope.organization(organization_id), + ) + return MagicMock( + sealed_secret=sealed.ciphertext, + kind=SecretKind.API_KEY.value, + key_version=sealed.key_version, + purpose="cohere", + ) + + +async def _resolve(kb, secret_row=None): + """Run the resolver against one KB row and an optional vault row. + + Returns the resolution and the secret-repo mock, so a test can assert the + vault was - or was not - consulted. + """ + with ( + patch(f"{_MODULE}.get_db_context") as db_ctx, + patch(f"{_MODULE}.knowledge_base_repo") as kbs, + patch(f"{_MODULE}.organization_secret_repo") as secrets, + ): + db_ctx.return_value.__aenter__ = AsyncMock(return_value=MagicMock()) + db_ctx.return_value.__aexit__ = AsyncMock(return_value=False) + kbs.get_for_collection = AsyncMock(return_value=kb) + secrets.get = AsyncMock(return_value=secret_row) + return await reranker_for_collection("handbook", organization_id=uuid.uuid4()), secrets + + +class TestResolution: + async def test_a_collection_nobody_claims_has_no_reranker(self): + resolved, secrets = await _resolve(None) + assert resolved is None + secrets.get.assert_not_called() + + async def test_a_collection_that_named_no_reranker_has_none(self): + """The normal off state - the vault is never consulted.""" + resolved, secrets = await _resolve(_kb(model=None, secret_id=None)) + assert resolved is None + secrets.get.assert_not_called() + + async def test_a_model_with_no_key_warns_that_a_configured_reranker_lost_its_key(self, caplog): + # Deleting the chosen secret nulls rerank_secret_id through the foreign + # key while leaving rerank_model set. That half state is a + # misconfiguration an operator should see, not the silent null/null off + # state - so it warns, and never reaches the vault (no key id to look up). + with caplog.at_level(logging.WARNING, logger=_MODULE): + resolved, secrets = await _resolve(_kb(model="rerank-v3.5", secret_id=None)) + assert resolved is None + secrets.get.assert_not_called() + assert "rerank_secret_missing" in caplog.text + + async def test_a_configured_collection_unseals_and_returns_its_reranker(self): + resolved, _ = await _resolve(_kb(secret_id=uuid.uuid4()), _sealed_key_row("co-org-own-key")) + assert resolved == ResolvedReranker(model="rerank-v3.5", api_key="co-org-own-key") + + async def test_a_personal_collection_never_looks_in_a_vault(self): + """No organization, no vault scope to open an envelope with.""" + resolved, secrets = await _resolve(_kb(secret_id=uuid.uuid4(), organization_id=None)) + assert resolved is None + secrets.get.assert_not_called() + + async def test_a_repr_never_carries_the_key(self): + resolved = ResolvedReranker(model="rerank-v3.5", api_key="co-secret") + assert "co-secret" not in repr(resolved) + assert "rerank-v3.5" in repr(resolved) + + async def test_the_secret_is_only_ever_looked_up_within_the_collections_org(self): + """The tenant boundary at the resolution layer: a collection's key is + fetched scoped to that collection's organization, so one organization's + collection can never resolve another's rerank key - it would not be + returned by the scoped query, and the vault would refuse the envelope + under the wrong scope (test_vault.py).""" + _, secrets = await _resolve(_kb(secret_id=uuid.uuid4()), _sealed_key_row("co-org-own-key")) + assert secrets.get.await_args.kwargs["organization_id"] == _ORG + + +class TestAuthorizedKnowledgeBaseId: + """A given `knowledge_base_id` resolves that exact row, not a name lookup. + + The search path passes the knowledge base access already authorized, so a + shared collection name cannot re-select a different same-named row and unseal + its key (#913).""" + + async def test_a_given_kb_id_reads_that_row_and_skips_the_name_lookup(self): + with ( + patch(f"{_MODULE}.get_db_context") as db_ctx, + patch(f"{_MODULE}.knowledge_base_repo") as kbs, + patch(f"{_MODULE}.organization_secret_repo") as secrets, + ): + db_ctx.return_value.__aenter__ = AsyncMock(return_value=MagicMock()) + db_ctx.return_value.__aexit__ = AsyncMock(return_value=False) + kbs.get_by_id = AsyncMock(return_value=_kb(model=None, secret_id=None)) + kbs.get_for_collection = AsyncMock(return_value=_kb(secret_id=uuid.uuid4())) + secrets.get = AsyncMock() + kb_id = uuid.uuid4() + + resolved = await reranker_for_collection( + "handbook", organization_id=_ORG, knowledge_base_id=kb_id + ) + + kbs.get_by_id.assert_awaited_once_with(kbs.get_by_id.await_args.args[0], kb_id) + kbs.get_for_collection.assert_not_called() + assert resolved is None + + +class TestDegradationTurnsRerankingOff: + """A chosen key that is gone drops reranking to off, with a line saying why. + + Never a 500: whose key pays for reranking must not decide whether a search + answers at all. + """ + + async def test_a_deleted_secret_turns_reranking_off_with_a_warning(self, caplog): + with caplog.at_level(logging.WARNING, logger=_MODULE): + resolved, _ = await _resolve(_kb(secret_id=uuid.uuid4()), None) + assert resolved is None + assert "rerank_secret_missing" in caplog.text + + async def test_an_unopenable_ciphertext_turns_reranking_off(self, caplog): + broken = MagicMock( + sealed_secret="not-a-ciphertext", kind=SecretKind.API_KEY.value, key_version=1 + ) + with caplog.at_level(logging.WARNING, logger=_MODULE): + resolved, _ = await _resolve(_kb(secret_id=uuid.uuid4()), broken) + assert resolved is None + assert "rerank_secret_unusable" in caplog.text + + async def test_a_secret_of_the_wrong_kind_turns_reranking_off(self, caplog): + row = _sealed_key_row("co-org") + with ( + patch(f"{_MODULE}.unseal_secret", return_value=MagicMock(spec=[])), + caplog.at_level(logging.WARNING, logger=_MODULE), + ): + resolved, _ = await _resolve(_kb(secret_id=uuid.uuid4()), row) + assert resolved is None + assert "rerank_secret_wrong_kind" in caplog.text + + async def test_the_normal_off_state_logs_nothing(self, caplog): + with caplog.at_level(logging.WARNING, logger=_MODULE): + await _resolve(_kb(model=None, secret_id=None)) + assert caplog.text == "" + + +class TestWhichReasonsAreDegraded: + def test_only_a_key_asked_for_and_not_given_is_degraded(self): + degraded = {source for source in RerankKeySource if source.is_degraded} + assert degraded == { + RerankKeySource.SECRET_MISSING, + RerankKeySource.SECRET_UNUSABLE, + RerankKeySource.SECRET_WRONG_KIND, + } diff --git a/backend/tests/test_reranker.py b/backend/tests/test_reranker.py new file mode 100644 index 000000000..88b4e0a02 --- /dev/null +++ b/backend/tests/test_reranker.py @@ -0,0 +1,255 @@ +"""Tests for the RAG reranker. + +Two properties carry the weight: the reranker reorders and re-scores the +candidates it is given, and its per-search cost is booked to whatever ledger is +metering the search - priced, because it is computed here rather than looked up +in a price table that does not know rerank models. +""" + +from __future__ import annotations + +from decimal import Decimal +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from app.agents.capabilities.budget import SpendLedger, metered_by +from app.services.rag.models import SearchResult +from app.services.rag.reranker import CohereReranker, build_reranker +from app.services.rerank_resolution import ResolvedReranker + +pytestmark = pytest.mark.anyio + + +def _results(*contents: str) -> list[SearchResult]: + return [SearchResult(content=c, score=0.0, metadata={"i": i}) for i, c in enumerate(contents)] + + +def _client_returning(*ranked: tuple[int, float], search_units: float | None = None) -> AsyncMock: + """A fake Cohere client whose rerank returns these (index, score) items. + + With `search_units` set, the response carries a `meta.billed_units` the way + Cohere's does; left None, the response has no `meta`, exercising the + candidate-count fallback. + """ + response = SimpleNamespace( + results=[SimpleNamespace(index=index, relevance_score=score) for index, score in ranked] + ) + if search_units is not None: + response.meta = SimpleNamespace(billed_units=SimpleNamespace(search_units=search_units)) + client = AsyncMock() + client.rerank = AsyncMock(return_value=response) + return client + + +def _reranker(client: AsyncMock) -> CohereReranker: + return CohereReranker(model="rerank-v3.5", api_key="co-key", client=client) + + +class TestReordering: + async def test_it_returns_the_candidates_in_the_rerankers_order(self): + client = _client_returning((2, 0.9), (0, 0.5), (1, 0.1)) + results = _results("first", "second", "third") + + reranked = await _reranker(client).rerank("q", results, top_n=3) + + assert [r.content for r in reranked] == ["third", "first", "second"] + + async def test_the_returned_score_is_the_rerankers_not_the_distance(self): + client = _client_returning((0, 0.87)) + reranked = await _reranker(client).rerank("q", _results("only"), top_n=1) + + assert reranked[0].score == 0.87 + assert reranked[0].metadata == {"i": 0} + + async def test_it_asks_cohere_for_at_most_the_candidates_it_has(self): + client = _client_returning((0, 0.9)) + await _reranker(client).rerank("q", _results("only"), top_n=5) + + assert client.rerank.await_args.kwargs["top_n"] == 1 + assert client.rerank.await_args.kwargs["documents"] == ["only"] + + async def test_no_candidates_returns_nothing_and_never_calls_cohere(self): + client = _client_returning() + reranked = await _reranker(client).rerank("q", [], top_n=5) + + assert reranked == [] + client.rerank.assert_not_awaited() + + +class TestSpend: + async def test_a_rerank_books_a_priced_nonzero_cost_to_the_active_ledger(self): + client = _client_returning((0, 0.9)) + ledger = SpendLedger() + + with metered_by(ledger): + await _reranker(client).rerank("q", _results("a", "b"), top_n=2) + + assert len(ledger.entries) == 1 + assert ledger.entries[0].priced + assert ledger.total_usd == Decimal("0.002") + + async def test_more_than_one_hundred_documents_bills_more_than_one_search_unit(self): + """The candidate-count fallback, used when the response omits billed_units.""" + client = _client_returning((0, 0.9)) + ledger = SpendLedger() + + with metered_by(ledger): + await _reranker(client).rerank("q", _results(*(str(n) for n in range(250))), top_n=5) + + assert ledger.total_usd == Decimal("0.006") + + async def test_it_meters_the_search_units_cohere_actually_billed(self): + """A few large chunks can bill several units; the response says how many.""" + client = _client_returning((0, 0.9), (1, 0.5), search_units=3) + ledger = SpendLedger() + + with metered_by(ledger): + await _reranker(client).rerank("q", _results("a", "b"), top_n=2) + + assert ledger.total_usd == Decimal("0.006") + + async def test_a_fractional_billed_unit_rounds_up(self): + client = _client_returning((0, 0.9), search_units=1.2) + ledger = SpendLedger() + + with metered_by(ledger): + await _reranker(client).rerank("q", _results("a"), top_n=1) + + assert ledger.total_usd == Decimal("0.004") + + async def test_a_failed_call_books_nothing_and_propagates(self): + client = AsyncMock() + client.rerank = AsyncMock(side_effect=RuntimeError("cohere down")) + ledger = SpendLedger() + + with metered_by(ledger), pytest.raises(RuntimeError): + await _reranker(client).rerank("q", _results("a"), top_n=1) + + assert ledger.entries == [] + + +class TestClientConstruction: + def test_the_client_is_built_lazily_from_the_key(self): + """No key is used and no client is built until the first rerank.""" + reranker = CohereReranker(model="rerank-v3.5", api_key="co-key") + assert reranker._client is None + assert reranker.client is not None + + +class TestClientLifecycle: + """A client the reranker builds is request-scoped and closed after the search. + + The process-wide retrieval service builds a fresh reranker per query, so an + unclosed client would leak one httpx connection pool per search. An injected + client is the caller's and is left open. + """ + + async def test_a_client_it_builds_is_closed_after_reranking(self): + built = _client_returning((0, 0.9)) + with patch("app.services.rag.reranker.cohere.AsyncClientV2", return_value=built) as ctor: + reranker = CohereReranker(model="rerank-v3.5", api_key="co-key") + await reranker.rerank("q", _results("a"), top_n=1) + + ctor.assert_called_once_with(api_key="co-key") + built.__aexit__.assert_awaited_once() + assert reranker._client is None + + async def test_a_built_client_is_closed_even_when_the_call_fails(self): + built = AsyncMock() + built.rerank = AsyncMock(side_effect=RuntimeError("cohere down")) + with patch("app.services.rag.reranker.cohere.AsyncClientV2", return_value=built): + reranker = CohereReranker(model="rerank-v3.5", api_key="co-key") + with pytest.raises(RuntimeError): + await reranker.rerank("q", _results("a"), top_n=1) + + built.__aexit__.assert_awaited_once() + assert reranker._client is None + + async def test_an_injected_client_is_left_open(self): + client = _client_returning((0, 0.9)) + reranker = _reranker(client) + + await reranker.rerank("q", _results("a"), top_n=1) + + client.__aexit__.assert_not_awaited() + assert reranker._client is client + + async def test_a_close_failure_is_swallowed_and_does_not_mask_the_result(self): + built = _client_returning((0, 0.9)) + built.__aexit__ = AsyncMock(side_effect=RuntimeError("pool already gone")) + with patch("app.services.rag.reranker.cohere.AsyncClientV2", return_value=built): + reranker = CohereReranker(model="rerank-v3.5", api_key="co-key") + reranked = await reranker.rerank("q", _results("a"), top_n=1) + + assert [r.content for r in reranked] == ["a"] + assert reranker._client is None + + +class TestConfigEquality: + """Two rerankers are equal when they name the same model and key. + + Retrieval leans on this to decide whether a multi-collection union shares + one reranker; the client is irrelevant to the comparison. + """ + + def test_same_model_and_key_are_equal(self): + assert CohereReranker("rerank-v3.5", "k") == CohereReranker("rerank-v3.5", "k") + + def test_the_client_does_not_affect_equality(self): + a = CohereReranker("rerank-v3.5", "k", client=AsyncMock()) + b = CohereReranker("rerank-v3.5", "k") + assert a == b + assert hash(a) == hash(b) + + def test_a_different_key_is_not_equal(self): + assert CohereReranker("rerank-v3.5", "k1") != CohereReranker("rerank-v3.5", "k2") + + def test_a_non_reranker_is_not_equal(self): + assert CohereReranker("rerank-v3.5", "k") != object() + + +class TestBuildReranker: + """The one composition point both retrieval paths share.""" + + async def test_an_unconfigured_collection_gets_no_reranker(self): + with patch( + "app.services.rag.reranker.reranker_for_collection", + new=AsyncMock(return_value=None), + ): + assert await build_reranker("handbook", None) is None + + async def test_a_configured_collection_gets_a_cohere_reranker(self): + with patch( + "app.services.rag.reranker.reranker_for_collection", + new=AsyncMock(return_value=ResolvedReranker(model="rerank-v3.5", api_key="co-key")), + ): + reranker = await build_reranker("handbook", None) + assert isinstance(reranker, CohereReranker) + assert reranker.model == "rerank-v3.5" + + +class TestBothPathsRerankThroughOneResolver: + """Done-when #3: spend is recorded on the agent-run path AND /rag/search. + + Both build their RetrievalService with the same `build_reranker`, so an + agent's knowledge search reranks exactly as the route does. Before this the + agent-run path built a RetrievalService with no resolver and never reranked. + """ + + def test_the_request_route_wires_build_reranker(self): + from app.api.deps import get_retrieval_service + + service = get_retrieval_service(MagicMock()) + assert service._reranker_resolver is build_reranker + + def test_the_agent_run_knowledge_tool_wires_build_reranker(self): + from app.agents.capabilities.knowledge import _search + + _search._retrieval_service = None + try: + service = _search.get_retrieval_service() + assert service._reranker_resolver is build_reranker + finally: + _search._retrieval_service = None diff --git a/backend/tests/test_reserved_collection_names.py b/backend/tests/test_reserved_collection_names.py index 27ab03e84..88730dab8 100644 --- a/backend/tests/test_reserved_collection_names.py +++ b/backend/tests/test_reserved_collection_names.py @@ -126,7 +126,7 @@ async def test_creating_a_collection_named_after_a_model_table_is_refused() -> N store, executed = _store() with pytest.raises(BadRequestError) as refused: - await store.create_collection("documents") + await store.create_collection("documents", organization_id=None) assert refused.value.details == {"collection": "documents", "table": "rag_documents"} assert executed == [] diff --git a/backend/tests/test_retrieval_reranking.py b/backend/tests/test_retrieval_reranking.py new file mode 100644 index 000000000..a15256c31 --- /dev/null +++ b/backend/tests/test_retrieval_reranking.py @@ -0,0 +1,219 @@ +"""How reranking changes what retrieval returns. + +The reranker is injected as a resolver, so these drive it with a stub that +reorders a known candidate set. Two properties matter: with a reranker, +retrieval returns the reranked order truncated to the limit; without one, it is +byte-for-byte the by-distance path. And a multi-collection search reranks the +union of candidates once, not each collection. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +import pytest + +from app.services.rag.models import SearchResult +from app.services.rag.reranker import BaseReranker, CohereReranker +from app.services.rag.retrieval import RetrievalService + +pytestmark = pytest.mark.anyio + + +class _ReverseReranker(BaseReranker): + """Reorders candidates worst-first, so a test can see reranking happened. + + Counts its calls, so a multi-collection search can assert it ran once over + the union rather than once per collection. + """ + + def __init__(self) -> None: + self.calls = 0 + + async def rerank( + self, query: str, results: list[SearchResult], top_n: int + ) -> list[SearchResult]: + self.calls += 1 + return list(reversed(results))[:top_n] + + +class _FailingReranker(BaseReranker): + async def rerank( + self, query: str, results: list[SearchResult], top_n: int + ) -> list[SearchResult]: + raise RuntimeError("cohere down") + + +def _store_returning(results: list[SearchResult]) -> MagicMock: + store = MagicMock() + store.search = AsyncMock(return_value=results) + return store + + +def _service(store: MagicMock, reranker: BaseReranker | None) -> RetrievalService: + settings = MagicMock() + settings.enable_hybrid_search = False + resolver = AsyncMock(return_value=reranker) if reranker is not None else None + return RetrievalService(vector_store=store, settings=settings, reranker_resolver=resolver) + + +def _service_by_name(store: MagicMock, mapping: dict[str, BaseReranker | None]) -> RetrievalService: + """A retrieval service whose reranker depends on which collection is asked.""" + settings = MagicMock() + settings.enable_hybrid_search = False + + async def resolver( + name: str, organization_id: object = None, knowledge_base_id: object = None + ) -> BaseReranker | None: + return mapping.get(name) + + return RetrievalService(vector_store=store, settings=settings, reranker_resolver=resolver) + + +def _hits(*names: str) -> list[SearchResult]: + return [SearchResult(content=name, score=score) for score, name in enumerate(names)] + + +class TestSingleCollection: + async def test_a_configured_collection_returns_the_reranked_order(self): + store = _store_returning(_hits("a", "b", "c")) + results = await _service(store, _ReverseReranker()).retrieve( + "q", "kb", limit=3, organization_id=None + ) + assert [r.content for r in results] == ["c", "b", "a"] + + async def test_it_truncates_to_the_limit_after_reranking(self): + store = _store_returning(_hits("a", "b", "c", "d")) + results = await _service(store, _ReverseReranker()).retrieve( + "q", "kb", limit=2, organization_id=None + ) + assert [r.content for r in results] == ["d", "c"] + + async def test_without_a_reranker_the_order_is_left_as_the_store_gave_it(self): + store = _store_returning(_hits("a", "b", "c")) + results = await _service(store, None).retrieve("q", "kb", limit=3, organization_id=None) + assert [r.content for r in results] == ["a", "b", "c"] + + async def test_a_reranker_failure_falls_back_to_the_distance_order(self): + store = _store_returning(_hits("a", "b", "c")) + results = await _service(store, _FailingReranker()).retrieve( + "q", "kb", limit=2, organization_id=None + ) + assert [r.content for r in results] == ["a", "b"] + + +class TestMultiCollection: + async def test_it_reranks_the_union_once_not_each_collection(self): + store = _store_returning(_hits("a", "b")) + reranker = _ReverseReranker() + await _service(store, reranker).retrieve_multi( + "q", collection_names=["kb_a", "kb_b"], limit=3, organization_id=None + ) + assert reranker.calls == 1 + + async def test_the_collection_stamp_survives_reranking(self): + store = _store_returning(_hits("a")) + results = await _service(store, _ReverseReranker()).retrieve( + "q", "handbook", limit=1, organization_id=None + ) + assert results[0].metadata["collection"] == "handbook" + + +class TestAuthorizedKnowledgeBaseIsPinned: + """The authorized KB id reaches resolution, so a shared collection name + resolves the row access granted rather than one re-looked-up by name (#913).""" + + @staticmethod + def _svc(store: MagicMock, resolver: AsyncMock) -> RetrievalService: + settings = MagicMock() + settings.enable_hybrid_search = False + return RetrievalService(vector_store=store, settings=settings, reranker_resolver=resolver) + + async def test_retrieve_threads_the_kb_id_to_the_resolver_and_the_store(self): + store = _store_returning(_hits("a")) + resolver = AsyncMock(return_value=None) + kb = uuid4() + + await self._svc(store, resolver).retrieve( + "q", "handbook", limit=1, organization_id=None, knowledge_base_id=kb + ) + + assert resolver.await_args.args == ("handbook", None, kb) + assert store.search.await_args.kwargs["knowledge_base_id"] == kb + + async def test_retrieve_multi_pins_each_collection_to_its_own_kb_id(self): + store = _store_returning(_hits("a")) + resolver = AsyncMock(return_value=None) + a, b = uuid4(), uuid4() + + await self._svc(store, resolver).retrieve_multi( + "q", + collection_names=["kb_a", "kb_b"], + limit=1, + organization_id=None, + knowledge_base_ids=[a, b], + ) + + resolved = {call.args[0]: call.args[2] for call in resolver.await_args_list} + assert resolved == {"kb_a": a, "kb_b": b} + + async def test_no_kb_id_falls_back_to_the_organization_scope(self): + store = _store_returning(_hits("a")) + resolver = AsyncMock(return_value=None) + + await self._svc(store, resolver).retrieve("q", "handbook", limit=1, organization_id=None) + + assert resolver.await_args.args == ("handbook", None, None) + + +class TestMixedRerankConfig: + """A union is reranked only when every collection agrees on one reranker. + + A set whose collections disagree - one reranking, one not, or two on + different keys - is left in distance order rather than reranked on a + credential that is not the collection's own. + """ + + async def test_a_differently_keyed_set_is_not_reranked(self): + store = _store_returning(_hits("a", "b")) + r1, r2 = _ReverseReranker(), _ReverseReranker() + svc = _service_by_name(store, {"kb_a": r1, "kb_b": r2}) + await svc.retrieve_multi( + "q", collection_names=["kb_a", "kb_b"], limit=3, organization_id=None + ) + assert r1.calls == 0 + assert r2.calls == 0 + + async def test_one_unconfigured_collection_disables_reranking_for_the_union(self): + store = _store_returning(_hits("a", "b")) + r = _ReverseReranker() + svc = _service_by_name(store, {"kb_a": r, "kb_b": None}) + await svc.retrieve_multi( + "q", collection_names=["kb_a", "kb_b"], limit=3, organization_id=None + ) + assert r.calls == 0 + + +class TestSharedReranker: + """`_shared_reranker` decides whether one reranker may reorder the union.""" + + def test_distinct_objects_with_one_config_are_shared(self): + a = CohereReranker("rerank-v3.5", "k") + b = CohereReranker("rerank-v3.5", "k") + assert RetrievalService._shared_reranker([a, b]) is a + + def test_a_differing_key_is_not_shared(self): + a = CohereReranker("rerank-v3.5", "k1") + b = CohereReranker("rerank-v3.5", "k2") + assert RetrievalService._shared_reranker([a, b]) is None + + def test_an_unconfigured_collection_in_the_set_disables_it(self): + a = CohereReranker("rerank-v3.5", "k") + assert RetrievalService._shared_reranker([a, None]) is None + + def test_an_all_unconfigured_set_has_no_reranker(self): + assert RetrievalService._shared_reranker([None, None]) is None + + def test_an_empty_set_has_no_reranker(self): + assert RetrievalService._shared_reranker([]) is None diff --git a/backend/tests/test_secrets.py b/backend/tests/test_secrets.py index 9614c8c8a..6ff31ccf0 100644 --- a/backend/tests/test_secrets.py +++ b/backend/tests/test_secrets.py @@ -34,6 +34,7 @@ ) from app.core.vault import VaultScope, seal from app.repositories import member_repo +from app.schemas.secret import SecretUsage from app.services.organization_secret import OrganizationSecretService from tests.test_model_profiles import service_account_json @@ -338,12 +339,63 @@ async def test_the_listing_carries_the_authors_id_for_a_stable_avatar_colour(sel "app.services.organization_secret.organization_secret_repo.agents_using", new=AsyncMock(return_value=[]), ), + patch( + "app.services.organization_secret.knowledge_base_repo.knowledge_bases_using", + new=AsyncMock(return_value=[]), + ), ): rows = await OrganizationSecretService(_db()).list_secrets(ctx) assert rows[0].created_by_user_id == author_id assert rows[0].created_by_email == "ada@acme.test" + @pytest.mark.anyio + async def test_a_key_bound_only_by_a_knowledge_base_is_not_reported_unused(self): + """A Cohere key used solely as a KB rerank (or embedding) credential still + shows what breaks on deletion. Reporting it unused - because only agent + specs were checked - invited deleting a key a collection was resolving, + which the SET NULL foreign key then silently turned off.""" + ctx = _ctx() + secret = _row(ctx, ApiKeySecret(api_key="co-live-4242")) + secret.description = None + secret.purpose = "custom" + secret.visibility = "org" + secret.owner_user_id = None + secret.created_by_user_id = None + secret.created_at = datetime.now(UTC) + secret.updated_at = None + kb_id = uuid.uuid4() + + with ( + patch( + "app.services.organization_secret.visible_resource_ids", + new=AsyncMock(return_value=None), + ), + patch( + "app.services.organization_secret.organization_secret_repo.list_secrets", + new=AsyncMock(return_value=[secret]), + ), + patch( + "app.services.organization_secret.member_repo.get_identities_for_users", + new=AsyncMock(return_value={}), + ), + patch( + "app.services.organization_secret.resource_grant_repo.count_for_resources", + new=AsyncMock(return_value={}), + ), + patch( + "app.services.organization_secret.organization_secret_repo.agents_using", + new=AsyncMock(return_value=[]), + ), + patch( + "app.services.organization_secret.knowledge_base_repo.knowledge_bases_using", + new=AsyncMock(return_value=[(kb_id, "Handbook")]), + ), + ): + rows = await OrganizationSecretService(_db()).list_secrets(ctx) + + assert rows[0].used_by == [SecretUsage(kind="knowledge_base", id=kb_id, name="Handbook")] + class TestRotatingASecret: @pytest.mark.anyio diff --git a/backend/tests/test_spend.py b/backend/tests/test_spend.py index c71ab6d8f..396d1465f 100644 --- a/backend/tests/test_spend.py +++ b/backend/tests/test_spend.py @@ -15,8 +15,10 @@ BudgetExceeded, BudgetGuard, BudgetScope, + SpendEntry, SpendLedger, SpendLimit, + book_ambient_spend, metered_by, price_request, record_ambient_usage, @@ -406,6 +408,46 @@ def test_nested_blocks_restore_the_outer_ledger(self): assert [entry.input_tokens for entry in outer.entries] == [2] +class TestBookAmbientSpend: + """A pre-priced entry booked to the active ledger without going through + `genai-prices`. + + The mechanism a reranker bills through: its cost comes from a published + per-search price, not from token counts, so it arrives already computed and + `priced=True` rather than being re-derived from a model name the price table + does not know. + """ + + def test_a_priced_entry_inside_a_metered_block_lands_priced(self): + ledger = SpendLedger() + entry = SpendEntry( + model_name="rerank-v3.5", + input_tokens=0, + output_tokens=0, + cost_usd=Decimal("0.002"), + priced=True, + ) + + with metered_by(ledger): + book_ambient_spend(entry) + + assert ledger.total_usd == Decimal("0.002") + assert not ledger.has_unpriced_models + + def test_a_priced_entry_with_nobody_metering_is_dropped_not_raised(self): + """A search outside any run has no ledger open; the reranker must still + rerank rather than refuse because nobody is counting.""" + book_ambient_spend( + SpendEntry( + model_name="rerank-v3.5", + input_tokens=0, + output_tokens=0, + cost_usd=Decimal("0.002"), + priced=True, + ) + ) + + class TestUsageDelta: """What a nested call added to a shared `RunUsage`, used by any capability that runs its own agent on `ctx.usage` - compaction, the LLM reminder.""" diff --git a/backend/tests/test_stats.py b/backend/tests/test_stats.py index f626d755f..5a39191a8 100644 --- a/backend/tests/test_stats.py +++ b/backend/tests/test_stats.py @@ -19,6 +19,7 @@ from app.core.exceptions import AuthorizationError, ValidationError from app.core.permissions import AuthContext, OrgRoleName +from app.db.models.ingestion_spend import SpendSource from app.services.stats import StatsService, resolve_window pytestmark = pytest.mark.anyio @@ -52,9 +53,20 @@ def repos(monkeypatch: pytest.MonkeyPatch) -> dict[str, AsyncMock]: mock = AsyncMock(return_value=value) monkeypatch.setattr(f"app.services.stats.agent_run_repo.{name}", mock) mocks[name] = mock + # One repo function serves both non-run sources, told apart by `source`, so + # the stub dispatches to a mock per source - each still called twice (this + # window, then the previous) and settable on its own. ingestion = AsyncMock(return_value=Decimal(0)) - monkeypatch.setattr("app.services.stats.ingestion_spend_repo.sum_cost_window", ingestion) + retrieval = AsyncMock(return_value=Decimal(0)) + + async def _sum_by_source(*args: object, source: object = None, **kwargs: object) -> Decimal: + if source == SpendSource.RETRIEVAL: + return await retrieval(*args, source=source, **kwargs) + return await ingestion(*args, source=source, **kwargs) + + monkeypatch.setattr("app.services.stats.ingestion_spend_repo.sum_cost_window", _sum_by_source) mocks["ingestion_sum_cost_window"] = ingestion + mocks["retrieval_sum_cost_window"] = retrieval member_count = AsyncMock(return_value=0) monkeypatch.setattr("app.services.stats.member_repo.count_for_org", member_count) mocks["count_for_org"] = member_count @@ -201,6 +213,24 @@ async def test_the_window_costs_models_plus_ingestion(self, repos) -> None: # bill against half of one. assert result.cost.previous_period_usd == Decimal("1.25") + async def test_retrieval_is_split_out_from_indexing_and_both_reach_the_total( + self, repos + ) -> None: + # A metered search shares the ingestion table but is reported apart, so + # it lands in retrieval_usd rather than inflating the indexing subtotal - + # while still counting toward the period total. + repos["sum_cost_window"].side_effect = [Decimal("2.00"), Decimal("1.00")] + repos["ingestion_sum_cost_window"].side_effect = [Decimal("0.50"), Decimal("0.25")] + repos["retrieval_sum_cost_window"].side_effect = [Decimal("0.10"), Decimal("0.05")] + + result = await StatsService(MagicMock()).usage(_ctx()) + + assert result.cost is not None + assert result.cost.ingestion_usd == Decimal("0.50") + assert result.cost.retrieval_usd == Decimal("0.10") + assert result.cost.period_usd == Decimal("2.60") + assert result.cost.previous_period_usd == Decimal("1.30") + async def test_own_scope_is_not_billed_for_the_organizations_indexing(self, repos) -> None: # `ingestion_spend` records no user - a document is indexed by a worker - # so charging a member's own window for a collection somebody else @@ -211,8 +241,10 @@ async def test_own_scope_is_not_billed_for_the_organizations_indexing(self, repo assert result.cost is not None assert result.cost.ingestion_usd == Decimal(0) + assert result.cost.retrieval_usd == Decimal(0) assert result.cost.period_usd == Decimal("2.00") repos["ingestion_sum_cost_window"].assert_not_called() + repos["retrieval_sum_cost_window"].assert_not_called() async def test_the_previous_total_is_asked_of_the_previous_window(self, repos) -> None: repos["count_runs"].side_effect = [40, 31] diff --git a/backend/tests/test_subagent_resolution.py b/backend/tests/test_subagent_resolution.py index bf90d25ed..f717a26b0 100644 --- a/backend/tests/test_subagent_resolution.py +++ b/backend/tests/test_subagent_resolution.py @@ -1004,6 +1004,7 @@ async def test_one_reaches_nothing_the_agent_that_invented_it_was_granted(self): ) assert built["resources"] == { "kb_collection_names": [], + "kb_collection_ids": [], "skills": [], "context_files": [], } diff --git a/docs/file-processing.md b/docs/file-processing.md index ca59d3f8d..3890357b0 100644 --- a/docs/file-processing.md +++ b/docs/file-processing.md @@ -468,6 +468,84 @@ the one caller that never asked the resolver at all, so every uploaded document was embedded with the deployment's model and key whatever its collection had chosen. +Both resolvers take the organization the ingest or search acts for, because the +vector namespace is deployment-global and a collection name is not unique across +tenants — two organizations may hold one name. Resolving by name alone returned +whichever row the database yielded first, so a shared name could embed or rerank +on another tenant's model and unseal its key. `get_for_collection` narrows the +candidates to the caller's own row, falling back to a deployment-wide +(`app`-scoped) collection and, only for a CLI ingest with no tenant, to the name +alone. + +Even own-org narrowing is not enough, though: an `app` collection everyone may +read and a restricted `org` collection of the same name are both the caller's to +resolve, but the caller may have been authorized for only the first. So the two +paths that search on a caller's behalf pass the **authorized** knowledge base's +id down to the resolvers, which read that exact row rather than looking one up by +name. On `POST /rag/search` the id comes from the same `readable_all` that +granted access; on an agent run it is the collection the agent's spec bound, +carried through `AgentDeps` beside its name. Either way resolution reads the row +the caller was actually granted, never a same-named one it selects first (#913). +Ingestion and the CLI, which choose the row themselves and have no distinct +authorized identity, pass none and keep the `organization_id`-scoped +`get_for_collection` lookup. + +### Reranking — a second pass, off unless configured + +Vector search orders results by embedding distance, which is a proxy for +relevance and sometimes a poor one. A **reranker** is an optional second pass: a +model scores each candidate against the query directly and reorders them, so a +better answer sitting well below the top by distance can surface. Retrieval +overfetches a wider candidate net (four times the limit rather than two), +reranks, then truncates — for a multi-collection search, once over the union of +every collection's candidates. That one reranker reorders the union only when +*every* collection resolves to the same one — same model, same key. An agent's +bound collections share one organization and one configuration, so the choice is +unambiguous there, but `POST /rag/search` may pass any readable set of one +organization; a set whose collections disagree — one reranking, one not, or two +on different keys — is left in distance order rather than reranked, so a +collection that opted out never has its candidates sent to Cohere on another +collection's credential. + +Configured **per knowledge base**, mirroring embeddings, by +`app/services/rerank_resolution.py`: + +| | | +|---|---| +| **Provider and model** | Cohere Rerank 3.5 is the first and only implementation behind `BaseReranker` (`app/services/rag/reranker.py`); a second provider is a second implementation, not a rewrite. `rerank_model` on the knowledge base names it, and unlike the embedding model it can be changed later. It is validated at create and update against `SUPPORTED_RERANK_MODELS` — a model this deployment cannot run would be stored and shown as configured, then swallowed as a failed Cohere call on every search, so an unsupported name is refused where the person setting it can see why. | +| **Credential** | The Cohere vault key chosen on the collection (`rerank_secret_id`), which is what the organization is billed for. | + +The one deliberate difference from embeddings is that reranking is **off by +default**. There is no deployment reranker key, so reranking runs only when a +collection sets *both* `rerank_model` and `rerank_secret_id`; either unset — or a +chosen key that is missing, unusable, or the wrong kind — resolves to no +reranker, and retrieval is byte-for-byte its pre-feature self. The three +misconfiguration cases are logged (a chosen key that vanished is an operator's +problem); the normal off state is silent. A runtime failure of the Cohere call +degrades to the by-distance order rather than failing the search — reranking is +an improvement on a working retrieval, not a dependency of it. The key is +validated at creation, the same as the embedding key and for the same reason. + +Rerank spend is metered like embeddings, with one deliberate exception. Ingestion +and the agent-run path already hold a ledger open, so a rerank during a knowledge +search books automatically — but `SpendLedger.record()` prices through +`genai-prices`, which is token-based and does not know rerank models, and would +book `cost_usd=0, priced=False`. So rerank does not go through `record()`: the +per-search Cohere cost is computed from a published per-search price +(`app/services/rag/reranker.py`, a constant checked against cohere.com/pricing +and dated in a comment) times the search units Cohere reports it billed — +`meta.billed_units.search_units` on the response, which counts the documents it +split past its token threshold, falling back to one unit per 100 candidates only +when the response omits the figure — and booked with `book_ambient_spend`, landing +`priced=True`. `POST /rag/search` — which opened no metering block and so left +even its embeddings unbilled (#16 class) — is now wrapped in one by +`KnowledgeSearchService`, so both its rerank and its embedding spend reach the +organization's monthly bill. And because a search is a paid call, that service +asserts the organization's monthly budget before it opens the ledger — the same +`assert_organization_within_budget` guard ingestion carries — so an exhausted cap +refuses the search before any embedding or Cohere call, rather than only being +recorded after the money is spent. + ### Vector Storage Vectors are stored in **pgvector** using the existing PostgreSQL database. No additional services needed. diff --git a/docs/governance.md b/docs/governance.md index 899f4a950..6484c5a00 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -352,18 +352,25 @@ which is `subagents-pydantic-ai` 0.2.20 and the reason the floor is there. **The dashboard's windowed figure carries it too.** `GET /stats/usage` answers a `cost` block for whatever period the filter chose, and that block is runs *plus* -ingestion — the same arithmetic the monthly cap is measured with — with -`model_usd` and `ingestion_usd` beside it so a reader can see where the money -went without subtracting. It reported the model half alone until 0.0.152, which -put two different definitions of cost on one card: the headline moved with the -period filter and counted runs, while the month-to-date line under it counted -the whole bill, and nothing said they were answering different questions. On a -deployment that indexes documents they simply disagreed. - -At `scope=own` the ingestion half is zero rather than a share: a document is -indexed by a worker and `ingestion_spend` records no user, so charging one -person's window for a collection somebody else synced would be inventing their -spend. +ingestion *plus* retrieval — the same arithmetic the monthly cap is measured with +— with `model_usd`, `ingestion_usd` and `retrieval_usd` beside it so a reader can +see where the money went without subtracting. It reported the model half alone +until 0.0.152, which put two different definitions of cost on one card: the +headline moved with the period filter and counted runs, while the month-to-date +line under it counted the whole bill, and nothing said they were answering +different questions. On a deployment that indexes documents they simply +disagreed. + +`retrieval_usd` is what a metered `POST /rag/search` spent on embeddings and +reranking. It shares the `ingestion_spend` table with indexing — both are RAG +spend on no agent run — but a `source` column keeps them apart, so a search is +reported as search rather than inflating the indexing subtotal. Both count +toward the monthly cap all the same. + +At `scope=own` the ingestion and retrieval halves are zero rather than a share: a +document is indexed by a worker and a colleague's search records no user of this +window, so charging one person's window for a collection somebody else synced or +searched would be inventing their spend. **Every query has to say which of the two it is answering**, and the first column is the default. The month-to-date figure and the per-agent breakdown behind it diff --git a/docs/secrets.md b/docs/secrets.md index e526d3a5b..862dd85f9 100644 --- a/docs/secrets.md +++ b/docs/secrets.md @@ -125,6 +125,14 @@ them through `vault.seal_fields`, which seals every field at one version and han that version back to store: the one way to write such a row, so "no version column" and "reset one field to v1" cannot be spelled by hand. +**Knowledge bases.** A collection may resolve its embedding key +(`embedding_secret_id`) and its rerank key (`rerank_secret_id`) from the +organization's vault. Both are `SET NULL` on delete, so deleting the key does not +break the row — it silently stops embedding or reranking. The vault's usage +listing (`used_by`, "what breaks if I delete this") reports these bindings +alongside the agents that hold a key, so a Cohere key a collection reranks with +does not read as unused and invite exactly that deletion. + **Third-party services.** A small catalog of services an organization may bring its own key for: diff --git a/frontend/e2e/kb-ingestion.spec.ts b/frontend/e2e/kb-ingestion.spec.ts index 4cba887e3..0860c1108 100644 --- a/frontend/e2e/kb-ingestion.spec.ts +++ b/frontend/e2e/kb-ingestion.spec.ts @@ -114,7 +114,9 @@ test.describe("Ingestion settings", () => { // And it can be changed afterwards. The API replaces the object wholesale, // so this is also the only assertion that the dialog sent the nine fields // nobody touched alongside the one that was. - await page.getByRole("button", { name: "Edit" }).click(); + // Scoped to this panel: the page also carries a Reranking section with its + // own Edit, so a page-wide "Edit" is now two buttons. + await (await howItReads(page)).getByRole("button", { name: "Edit" }).click(); const settings = page.getByRole("dialog"); await settings.getByLabel("Chunk size").fill("2048"); await settings.getByRole("button", { name: "Save" }).click(); diff --git a/frontend/messages/en.json b/frontend/messages/en.json index aaf1eb1f7..6f1bf6bd7 100644 --- a/frontend/messages/en.json +++ b/frontend/messages/en.json @@ -1569,7 +1569,9 @@ "title": "Spend", "unit": "in this period", "description": "What the period cost, and this calendar month to date. The month figure reconciles with the invoice, so it deliberately ignores the period filter.", - "split": "{models} on models · {ingestion} on indexing" + "splitModels": "{amount} on models", + "splitIndexing": "{amount} on indexing", + "splitSearch": "{amount} on search" }, "surfaces": { "empty": { @@ -2033,6 +2035,17 @@ "reasoning": "Reasoning", "remove": "Remove", "removeNamed": "Remove {name}?", + "rerank": "Reranking", + "rerankBilledTo": "billed to {key}", + "rerankHelp": "A reranker reorders search results by relevance. Off unless a key is set.", + "rerankKey": "Reranking key", + "rerankKeyConfigured": "a Cohere key", + "rerankKeyName": "Cohere (reranking)", + "rerankOff": "Off", + "rerankOffExplained": "Off. Results are ordered by vector distance alone.", + "rerankOnExplained": "A second pass reorders each search's results by relevance before they are returned.", + "rerankSettings": "Reranking", + "rerankSettingsDescription": "Choose which key reranks searches against {name}. It takes effect from the next search; nothing already stored changes.", "reusableIntegrations": "Reusable integrations", "save": "Save", "saving": "Saving…", @@ -2078,6 +2091,7 @@ "failedUpdate": "Failed to update knowledge base", "inactive": "Inactive", "ingestionSaved": "Ingestion settings saved", + "rerankSaved": "Reranking updated", "integrationCloned": "Integration cloned to this knowledge base", "namePlaceholder": "Product docs", "scopeApp": "App-wide", @@ -2724,6 +2738,10 @@ "title": "Embeddings (leave the default)", "body": "How documents are indexed for search. It's frozen at creation, but the deployment default is right for almost everyone — expand it only if you know you need a different model. Choose Next." }, + "flow-kb-field-rerank": { + "title": "Reranking (optional)", + "body": "A reranker reorders search results by relevance. It stays off unless you pick a key to pay for it, and the default off is right for almost everyone. Choose Next." + }, "flow-kb-field-create": { "title": "Create it", "body": "Choose Create. We'll move on the moment it exists — you'll add documents to it right after." @@ -2924,6 +2942,10 @@ "title": "How documents are read", "body": "The parser, chunking and embedding model behind this collection — what turns your files into something searchable. Change it, and new uploads follow the new settings." }, + "kb-rerank": { + "title": "Sharpen the results", + "body": "A reranker adds a second pass that reorders search results by relevance. Off by default; turn it on here with a Cohere key, or turn it back off — it changes from the next search on." + }, "kb-sync": { "title": "Keep it in step", "body": "Connect a drive, a site or a repo, and the collection refreshes itself from the source instead of waiting on a manual re-upload." diff --git a/frontend/messages/pl.json b/frontend/messages/pl.json index 658e209e3..31ab510aa 100644 --- a/frontend/messages/pl.json +++ b/frontend/messages/pl.json @@ -684,7 +684,9 @@ "description": "Koszt modeli ląduje tutaj uruchomienie po uruchomieniu." }, "description": "Ile kosztował okres i ile miesiąc kalendarzowy do dziś. Kwota miesięczna zgadza się z fakturą, więc celowo nie reaguje na filtr okresu.", - "split": "{models} na modele · {ingestion} na indeksowanie" + "splitModels": "{amount} na modele", + "splitIndexing": "{amount} na indeksowanie", + "splitSearch": "{amount} na wyszukiwanie" }, "model-mix": { "title": "Modele za uruchomieniami", diff --git a/frontend/src/app/[locale]/(dashboard)/rag/[id]/counts.integration.test.tsx b/frontend/src/app/[locale]/(dashboard)/rag/[id]/counts.integration.test.tsx index 7147f4ebd..1a8583a04 100644 --- a/frontend/src/app/[locale]/(dashboard)/rag/[id]/counts.integration.test.tsx +++ b/frontend/src/app/[locale]/(dashboard)/rag/[id]/counts.integration.test.tsx @@ -53,6 +53,8 @@ const COLLECTION: KnowledgeBase = { embedding_provider: "openrouter", embedding_secret_id: null, embedding_dim: 1536, + rerank_model: null, + rerank_secret_id: null, created_at: "2026-01-01T00:00:00Z", updated_at: null, // Zero, as the single-row read really answers - the three counts are derived diff --git a/frontend/src/app/[locale]/(dashboard)/rag/[id]/delete-collection.integration.test.tsx b/frontend/src/app/[locale]/(dashboard)/rag/[id]/delete-collection.integration.test.tsx index 5b3f23be6..92ad21ecc 100644 --- a/frontend/src/app/[locale]/(dashboard)/rag/[id]/delete-collection.integration.test.tsx +++ b/frontend/src/app/[locale]/(dashboard)/rag/[id]/delete-collection.integration.test.tsx @@ -64,6 +64,8 @@ const COLLECTION: KnowledgeBase = { embedding_provider: "openrouter", embedding_secret_id: null, embedding_dim: 1536, + rerank_model: null, + rerank_secret_id: null, created_at: "2026-01-01T00:00:00Z", updated_at: null, // Zero, as the single-row read really answers: the three counts are derived diff --git a/frontend/src/app/[locale]/(dashboard)/rag/[id]/kb-detail-sections.integration.test.tsx b/frontend/src/app/[locale]/(dashboard)/rag/[id]/kb-detail-sections.integration.test.tsx index eb83a89b3..c911d40c4 100644 --- a/frontend/src/app/[locale]/(dashboard)/rag/[id]/kb-detail-sections.integration.test.tsx +++ b/frontend/src/app/[locale]/(dashboard)/rag/[id]/kb-detail-sections.integration.test.tsx @@ -66,6 +66,8 @@ const KB: KnowledgeBase = { embedding_provider: "openrouter", embedding_secret_id: null, embedding_dim: 3072, + rerank_model: null, + rerank_secret_id: null, created_at: "2026-07-01T00:00:00Z", updated_at: null, document_count: 0, diff --git a/frontend/src/app/[locale]/(dashboard)/rag/[id]/page.tsx b/frontend/src/app/[locale]/(dashboard)/rag/[id]/page.tsx index c1156af20..a0ef44dc9 100644 --- a/frontend/src/app/[locale]/(dashboard)/rag/[id]/page.tsx +++ b/frontend/src/app/[locale]/(dashboard)/rag/[id]/page.tsx @@ -27,6 +27,8 @@ import { FileViewer } from "@/components/kb/file-viewer"; import { EmbeddingDialog } from "@/components/kb/embedding-dialog"; import { IngestionDialog } from "@/components/kb/ingestion-dialog"; import { IngestionPanel } from "@/components/kb/ingestion-panel"; +import { RerankDialog } from "@/components/kb/rerank-dialog"; +import { RerankPanel } from "@/components/kb/rerank-panel"; import { UploadOverrideDialog } from "@/components/kb/upload-override-dialog"; import { useKBDetail, usePermissions, usePollWhileIngesting, useUrlState } from "@/hooks"; import { overrideSize } from "@/lib/ingestion-config"; @@ -85,6 +87,7 @@ export default function KBDetailPage({ params }: KBDetailPageProps) { refresh, loadMoreDocuments, updateIngestion, + updateRerank, updateEmbeddings, uploadDocument, deleteDocument, @@ -105,6 +108,7 @@ export default function KBDetailPage({ params }: KBDetailPageProps) { const [creatingSource, setCreatingSource] = useState(false); const [viewerDoc, setViewerDoc] = useState(null); const [ingestionOpen, setIngestionOpen] = useState(false); + const [rerankOpen, setRerankOpen] = useState(false); const [embeddingOpen, setEmbeddingOpen] = useState(false); const [overrideOpen, setOverrideOpen] = useState(false); /** @@ -299,6 +303,16 @@ export default function KBDetailPage({ params }: KBDetailPageProps) { onEditEmbeddings={mayEdit ? () => setEmbeddingOpen(true) : undefined} /> + {/* Reranking is the other per-collection retrieval knob, and the only + one changeable after creation. No Edit on an app-scoped collection - + it carries no organization_id, so it can hold no vault key and the + backend would refuse one; the panel stays as a read-only fact. */} +
+ setRerankOpen(true) : undefined} + /> +
@@ -407,6 +421,14 @@ export default function KBDetailPage({ params }: KBDetailPageProps) { onSave={updateIngestion} /> + + = {}): KnowledgeBase { embedding_provider: "openrouter", embedding_secret_id: null, embedding_dim: 3072, + rerank_model: null, + rerank_secret_id: null, created_at: "2026-07-01T00:00:00Z", updated_at: null, document_count: 0, diff --git a/frontend/src/components/dashboard/widgets/spend.test.tsx b/frontend/src/components/dashboard/widgets/spend.test.tsx new file mode 100644 index 000000000..8df61494f --- /dev/null +++ b/frontend/src/components/dashboard/widgets/spend.test.tsx @@ -0,0 +1,95 @@ +import { render, screen } from "@testing-library/react"; +import { NextIntlClientProvider } from "next-intl"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import messages from "../../../../messages/en.json"; +import { SpendWidget } from "./spend"; +import type { Period } from "@/lib/dashboard/period"; +import type { CostBlock } from "@/types/stats"; + +/** + * The headline is the whole bill; the caption underneath splits it into the + * parts that spent. A part at zero is left off, so a deployment that only ran + * models reads no split at all - and a search is reported as search, never + * folded into indexing. + */ + +const useUsageStatsMock = vi.fn(); +const useSpendMock = vi.fn(); +vi.mock("@/hooks", () => ({ + useUsageStats: (...args: unknown[]) => useUsageStatsMock(...args), + useSpend: (...args: unknown[]) => useSpendMock(...args), +})); + +const PERIOD: Period = { preset: "30d", from: "2026-07-07", to: "2026-08-05" }; + +function withCost(cost: Partial) { + useUsageStatsMock.mockReturnValue({ + usage: { total_runs: 5, cost: { by_provider: [], ...cost } }, + isLoading: false, + isStale: false, + error: null, + refetch: vi.fn(), + }); + useSpendMock.mockReturnValue({ spend: null }); +} + +function renderWidget() { + return render( + + + , + ); +} + +beforeEach(() => { + useUsageStatsMock.mockReset(); + useSpendMock.mockReset(); +}); + +describe("the spend widget", () => { + it("splits the bill into models, indexing and search when each spent", () => { + withCost({ + period_usd: "2.60", + previous_period_usd: "1.30", + model_usd: "2.00", + ingestion_usd: "0.50", + retrieval_usd: "0.10", + }); + renderWidget(); + + expect(screen.getByText(/on models/)).toBeInTheDocument(); + expect(screen.getByText(/on indexing/)).toBeInTheDocument(); + expect(screen.getByText(/on search/)).toBeInTheDocument(); + }); + + it("shows no split when only model requests spent", () => { + withCost({ + period_usd: "2.00", + previous_period_usd: "1.00", + model_usd: "2.00", + ingestion_usd: "0", + retrieval_usd: "0", + }); + renderWidget(); + + expect(screen.queryByText(/on models/)).toBeNull(); + expect(screen.queryByText(/on indexing/)).toBeNull(); + expect(screen.queryByText(/on search/)).toBeNull(); + }); + + it("reports search without indexing when a search spent but nothing was indexed", () => { + withCost({ + period_usd: "2.10", + previous_period_usd: "1.00", + model_usd: "2.00", + ingestion_usd: "0", + retrieval_usd: "0.10", + }); + renderWidget(); + + expect(screen.getByText(/on models/)).toBeInTheDocument(); + expect(screen.queryByText(/on indexing/)).toBeNull(); + expect(screen.getByText(/on search/)).toBeInTheDocument(); + }); +}); diff --git a/frontend/src/components/dashboard/widgets/spend.tsx b/frontend/src/components/dashboard/widgets/spend.tsx index d340bf6ad..315ddeb88 100644 --- a/frontend/src/components/dashboard/widgets/spend.tsx +++ b/frontend/src/components/dashboard/widgets/spend.tsx @@ -38,6 +38,17 @@ export function SpendWidget({ title, hint, period, seeAll, options }: DashboardW const current = Number(cost?.period_usd ?? 0); const previous = Number(cost?.previous_period_usd ?? 0); const delta = deltaPercent(current, previous); + // The bill's parts, shown only where they spent: models always, then + // indexing and search when a knowledge base was used. A deployment + // with none reads no split at all, and the parts join rather than + // sitting in the provider bars below, which break down the model half. + const splitParts = [t("splitModels", { amount: formatUsd(cost?.model_usd) })]; + if (Number(cost?.ingestion_usd ?? 0) > 0) { + splitParts.push(t("splitIndexing", { amount: formatUsd(cost?.ingestion_usd) })); + } + if (Number(cost?.retrieval_usd ?? 0) > 0) { + splitParts.push(t("splitSearch", { amount: formatUsd(cost?.retrieval_usd) })); + } return (
) : undefined } - // The two halves of the bill, and only when indexing spent - // anything: a deployment with no knowledge base should not - // read a line about a subsystem it does not use. The bars - // below break down the model half, so the split rides the - // headline rather than joining them - two denominators in one - // list read as one. - caption={ - Number(cost?.ingestion_usd ?? 0) > 0 - ? t("split", { - models: formatUsd(cost?.model_usd), - ingestion: formatUsd(cost?.ingestion_usd), - }) - : undefined - } + // The parts of the bill, shown only once something beyond model + // spend was billed: a deployment with no knowledge base should + // not read a line about a subsystem it does not use. + caption={splitParts.length > 1 ? splitParts.join(" · ") : undefined} /> ({ diff --git a/frontend/src/components/kb/create-kb-dialog.integration.test.tsx b/frontend/src/components/kb/create-kb-dialog.integration.test.tsx index 491b5e043..610ef27b3 100644 --- a/frontend/src/components/kb/create-kb-dialog.integration.test.tsx +++ b/frontend/src/components/kb/create-kb-dialog.integration.test.tsx @@ -394,10 +394,11 @@ describe("the two keys this dialog can store", () => { expect(screen.queryByRole("button", { name: "Add a key: OpenRouter (embeddings)" })).toBeNull(); expect(screen.queryByRole("button", { name: "Add a key: OpenAI" })).toBeNull(); - // Two sentences, not silence, one per offer: the inline form says it here, - // and the model panel says it in its own words because a disabled Add model - // with nothing beside it explains nothing. - expect(screen.getAllByText(/permission you do not hold/)).toHaveLength(2); + expect(screen.queryByRole("button", { name: "Add a key: Cohere (reranking)" })).toBeNull(); + // A sentence, not silence, one per offer: the embedding and rerank inline + // forms say it here, and the model panel says it in its own words because a + // disabled Add model with nothing beside it explains nothing. + expect(screen.getAllByText(/permission you do not hold/)).toHaveLength(3); }); }); diff --git a/frontend/src/components/kb/create-kb-dialog.test.tsx b/frontend/src/components/kb/create-kb-dialog.test.tsx index d086cbd4f..cb657c2f9 100644 --- a/frontend/src/components/kb/create-kb-dialog.test.tsx +++ b/frontend/src/components/kb/create-kb-dialog.test.tsx @@ -60,6 +60,21 @@ describe("CreateKBDialog", () => { // there is a `TypeError` in `usePermissions`, not "no permissions". if (path === "/me/permissions") return { organization_id: "org-1", role: "member", is_app_admin: false, permissions: [] }; + // One key that can pay for reranking, so the picker has something to + // choose beyond Off. Its purpose is what makes it a rerank key. + if (path === "/secrets") + return { + items: [ + { + id: "cohere-1", + name: "Cohere key", + hint: "4242", + purpose: "cohere", + kind: "api_key", + }, + ], + total: 1, + }; return { items: [], total: 0 }; }); vi.mocked(apiClient.post).mockResolvedValue({ id: "kb-1", name: "Handbook" }); @@ -154,6 +169,31 @@ describe("CreateKBDialog", () => { ); }); + it("sends no reranking fields for a collection nobody turned it on for", async () => { + // Both-or-neither: the backend reads reranking as on only when the model and + // the key arrive together, so leaving it Off means sending neither. + await userEvent.type(screen.getByLabelText("Name"), "Handbook"); + await userEvent.click(create()); + + await waitFor(() => expect(apiClient.post).toHaveBeenCalled()); + expect(posted()).not.toHaveProperty("rerank_secret_id"); + expect(posted()).not.toHaveProperty("rerank_model"); + }); + + it("sends the key and the one model once a reranking key is chosen", async () => { + // The model is a constant, not a choice: there is one reranker and no + // endpoint listing them, so choosing the key is choosing to rerank. + await userEvent.type(screen.getByLabelText("Name"), "Handbook"); + await userEvent.click(screen.getByText("Reranking")); + await userEvent.click(screen.getByLabelText("Reranking key")); + await userEvent.click(await screen.findByRole("option", { name: /Cohere key/ })); + await userEvent.click(create()); + + await waitFor(() => expect(apiClient.post).toHaveBeenCalled()); + expect(posted().rerank_secret_id).toBe("cohere-1"); + expect(posted().rerank_model).toBe("rerank-v3.5"); + }); + it("keeps the name and its own settings when the server refuses", async () => { // A dialog that clears itself on a refusal makes the refusal cost the whole // form, which is how people learn not to open the settings at all. diff --git a/frontend/src/components/kb/create-kb-dialog.tsx b/frontend/src/components/kb/create-kb-dialog.tsx index caabc4552..a85a07442 100644 --- a/frontend/src/components/kb/create-kb-dialog.tsx +++ b/frontend/src/components/kb/create-kb-dialog.tsx @@ -24,8 +24,10 @@ import { } from "@/components/ui/select"; import { EmbeddingProviderFields, useEmbeddingProviders } from "@/components/kb/embedding-picker"; import { IngestionSettings } from "@/components/kb/ingestion-settings"; +import { InlineSecret } from "@/components/vault/inline-secret"; import { ProviderRow } from "@/components/vault/provider-row"; import { useKnowledgeBases } from "@/hooks"; +import { useSecrets } from "@/hooks/use-secrets"; import { submitFailure } from "@/lib/api-error"; import { DEFAULT_INGESTION_CONFIG, @@ -33,6 +35,7 @@ import { ingestionProblems, sameIngestion, } from "@/lib/ingestion-config"; +import { DEFAULT_RERANK_MODEL, RERANK_KEY_PURPOSE, RERANK_OFF } from "@/lib/rerank-config"; import { cn } from "@/lib/utils"; import type { CreateKnowledgeBaseInput, IngestionConfig, KBScope } from "@/types"; import { useTranslations } from "next-intl"; @@ -58,9 +61,12 @@ export function CreateKBDialog({ open, onOpenChange, onCreated }: CreateKBDialog const [embeddingModel, setEmbeddingModel] = useState(null); const [embeddingProvider, setEmbeddingProvider] = useState(null); const [embeddingSecretId, setEmbeddingSecretId] = useState(null); + const [rerankSecretId, setRerankSecretId] = useState(null); const [isSubmitting, setIsSubmitting] = useState(false); const [errors, setErrors] = useState>>({}); const { createKB } = useKnowledgeBases(); + const { secrets } = useSecrets(); + const rerankKeys = secrets.filter((secret) => secret.purpose === RERANK_KEY_PURPOSE); const { models: embeddingModels, unreadable: modelsUnreadable } = useEmbeddingProviders(); // Whose endpoint the models on offer belong to. The provider decides both the // model list and which vault keys can pay, so it is resolved before either - @@ -94,6 +100,7 @@ export function CreateKBDialog({ open, onOpenChange, onCreated }: CreateKBDialog setEmbeddingModel(null); setEmbeddingProvider(null); setEmbeddingSecretId(null); + setRerankSecretId(null); setErrors({}); }; @@ -115,6 +122,12 @@ export function CreateKBDialog({ open, onOpenChange, onCreated }: CreateKBDialog input.embedding_provider = provider; } if (embeddingSecretId) input.embedding_secret_id = embeddingSecretId; + // Both or neither: the backend turns reranking on only when the model and + // the key arrive together, so a key with no model would be a silent no-op. + if (rerankSecretId) { + input.rerank_secret_id = rerankSecretId; + input.rerank_model = DEFAULT_RERANK_MODEL; + } const kb = await createKB(input); reset(); onOpenChange(false); @@ -280,6 +293,48 @@ export function CreateKBDialog({ open, onOpenChange, onCreated }: CreateKBDialog
+
+ + + {t("rerank")} + + {rerankSecretId ? DEFAULT_RERANK_MODEL : t("rerankOff")} + + +
+

{t("rerankHelp")}

+
+ + + +
+
+
+ {/* Folded away, because creating a collection is a two-field job and most people will never open this. It is a disclosure rather than a diff --git a/frontend/src/components/kb/index.ts b/frontend/src/components/kb/index.ts index bf53017b9..c979d12e5 100644 --- a/frontend/src/components/kb/index.ts +++ b/frontend/src/components/kb/index.ts @@ -3,5 +3,7 @@ export { FileViewer } from "./file-viewer"; export { IngestionDialog } from "./ingestion-dialog"; export { IngestionPanel } from "./ingestion-panel"; export { IngestionSettings, type IngestionSettingsProps } from "./ingestion-settings"; +export { RerankDialog } from "./rerank-dialog"; +export { RerankPanel } from "./rerank-panel"; export { ReusableIntegrations } from "./reusable-integrations"; export { UploadOverrideDialog } from "./upload-override-dialog"; diff --git a/frontend/src/components/kb/rerank-dialog.test.tsx b/frontend/src/components/kb/rerank-dialog.test.tsx new file mode 100644 index 000000000..509c28c9b --- /dev/null +++ b/frontend/src/components/kb/rerank-dialog.test.tsx @@ -0,0 +1,227 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import type { ReactNode } from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { RerankDialog } from "./rerank-dialog"; +import { apiClient, ApiError } from "@/lib/api-client"; + +/** + * The edit dialog is the whole point of "reranking can be changed after + * creation": the create dialog sets it once, this turns it on, swaps its key or + * turns it off. What it must get right is the pair it sends - a key and the one + * model, or two nulls - because that pair is how the backend tells "change + * reranking" from "leave it alone". + */ + +vi.mock("@/lib/api-client", async (importOriginal) => ({ + ...(await importOriginal()), + apiClient: { get: vi.fn(), post: vi.fn(), patch: vi.fn(), delete: vi.fn() }, +})); +const toastError = vi.fn(); +vi.mock("sonner", () => ({ toast: { success: vi.fn(), error: (m: string) => toastError(m) } })); + +const SECRETS = { + items: [ + { id: "co-1", name: "Cohere prod", hint: "4242", purpose: "cohere", kind: "api_key" }, + // Not a rerank key: it must never be offered as one that can pay for reranking. + { id: "tav-1", name: "Tavily", hint: "9999", purpose: "tavily", kind: "api_key" }, + ], + total: 2, +}; + +function wrapper({ children }: { children: ReactNode }) { + const client = new QueryClient({ + defaultOptions: { queries: { retry: false }, mutations: { retry: false } }, + }); + return {children}; +} + +beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(apiClient.get).mockImplementation(async (path: string) => { + if (path === "/secrets") return SECRETS; + if (path === "/me/permissions") + return { organization_id: "org-1", role: "builder", is_app_admin: false, permissions: [] }; + return { items: [], total: 0 }; + }); +}); + +describe("choosing a key", () => { + it("offers only keys that can pay for reranking", async () => { + render( + , + { wrapper }, + ); + await userEvent.click(screen.getByLabelText("Reranking key")); + + expect(await screen.findByRole("option", { name: /Cohere prod/ })).toBeInTheDocument(); + expect(screen.queryByRole("option", { name: /Tavily/ })).toBeNull(); + }); + + it("turning it on sends the one model paired with the chosen key", async () => { + const onSave = vi.fn().mockResolvedValue({}); + const onOpenChange = vi.fn(); + render( + , + { wrapper }, + ); + + await userEvent.click(screen.getByLabelText("Reranking key")); + await userEvent.click(await screen.findByRole("option", { name: /Cohere prod/ })); + await userEvent.click(screen.getByRole("button", { name: "Save" })); + + expect(onSave).toHaveBeenCalledWith({ + rerank_model: "rerank-v3.5", + rerank_secret_id: "co-1", + }); + await waitFor(() => expect(onOpenChange).toHaveBeenCalledWith(false)); + }); + + it("turning it off sends the pair as two nulls", async () => { + const onSave = vi.fn().mockResolvedValue({}); + render( + , + { wrapper }, + ); + + await userEvent.click(screen.getByLabelText("Reranking key")); + await userEvent.click(await screen.findByRole("option", { name: "Off" })); + await userEvent.click(screen.getByRole("button", { name: "Save" })); + + expect(onSave).toHaveBeenCalledWith({ rerank_model: null, rerank_secret_id: null }); + }); + + it("cannot be saved until something changes", async () => { + render( + , + { wrapper }, + ); + + expect(screen.getByRole("button", { name: "Save" })).toBeDisabled(); + }); +}); + +describe("when the save is refused", () => { + it("shows a key the server named as wrong beside the picker", async () => { + const onSave = vi.fn().mockRejectedValue( + new ApiError(422, "Invalid", { + error: { + code: "VALIDATION_ERROR", + message: "Invalid", + details: { + fields: [ + { + field: "rerank_secret_id", + message: "That key is for tavily; reranking runs through Cohere.", + }, + ], + }, + }, + }), + ); + render( + , + { wrapper }, + ); + + await userEvent.click(screen.getByLabelText("Reranking key")); + await userEvent.click(await screen.findByRole("option", { name: /Cohere prod/ })); + await userEvent.click(screen.getByRole("button", { name: "Save" })); + + expect(await screen.findByText(/That key is for tavily/)).toBeInTheDocument(); + }); + + it("toasts a refusal that names no field", async () => { + const onSave = vi.fn().mockRejectedValue(new Error("boom")); + render( + , + { wrapper }, + ); + + await userEvent.click(screen.getByLabelText("Reranking key")); + await userEvent.click(await screen.findByRole("option", { name: /Cohere prod/ })); + await userEvent.click(screen.getByRole("button", { name: "Save" })); + + await waitFor(() => expect(toastError).toHaveBeenCalledWith("boom")); + }); +}); + +describe("what it shows when reopened", () => { + it("discards an abandoned pick, re-seeding from the server on reopen", async () => { + // Save is disabled only when the draft equals what the server holds, so it + // is the observable proof of a re-seed: if the abandoned pick survived, the + // draft would differ from the server and Save would be live. + const props = { + onOpenChange: vi.fn(), + rerankSecretId: "co-1" as string | null, + collectionName: "handbook_x", + onSave: vi.fn(), + }; + const { rerender } = render(, { wrapper }); + + await userEvent.click(screen.getByLabelText("Reranking key")); + await userEvent.click(await screen.findByRole("option", { name: "Off" })); + expect(screen.getByRole("button", { name: "Save" })).toBeEnabled(); + + // Closed, then reopened against the unchanged collection: the draft goes back + // to its key and Save falls dormant again. + rerender(); + rerender(); + + expect(screen.getByRole("button", { name: "Save" })).toBeDisabled(); + }); + + it("follows the server's value when reranking is changed elsewhere while open", async () => { + // The pair moving under an open dialog (a save in another tab) re-seeds the + // draft too - Save stays dormant because the draft tracked the change rather + // than staying on the key that is no longer set. + const props = { + open: true, + onOpenChange: vi.fn(), + collectionName: "handbook_x", + onSave: vi.fn(), + }; + const { rerender } = render(, { wrapper }); + rerender(); + + expect(screen.getByRole("button", { name: "Save" })).toBeDisabled(); + }); +}); diff --git a/frontend/src/components/kb/rerank-dialog.tsx b/frontend/src/components/kb/rerank-dialog.tsx new file mode 100644 index 000000000..707c895df --- /dev/null +++ b/frontend/src/components/kb/rerank-dialog.tsx @@ -0,0 +1,155 @@ +"use client"; + +import { useState } from "react"; +import { toast } from "sonner"; + +import { + Button, + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, + Label, + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from "@/components/ui"; +import { InlineSecret } from "@/components/vault/inline-secret"; +import { ProviderRow } from "@/components/vault/provider-row"; +import { useSecrets } from "@/hooks"; +import { useChanged } from "@/hooks/use-changed"; +import { submitFailure } from "@/lib/api-error"; +import { DEFAULT_RERANK_MODEL, RERANK_KEY_PURPOSE, RERANK_OFF } from "@/lib/rerank-config"; +import type { UpdateRerankInput } from "@/types"; +import { useTranslations } from "next-intl"; + +interface RerankDialogProps { + open: boolean; + onOpenChange: (open: boolean) => void; + /** The rerank key the collection is set to now, or null when reranking is off. */ + rerankSecretId: string | null; + collectionName: string; + onSave: (input: UpdateRerankInput) => Promise; +} + +/** + * Turning reranking on, changing its key, or turning it off on an existing + * collection. + * + * A dialog rather than an inline control for the same reason the ingestion one + * is: the change is a decision with a cost attached (a rerank key is billed per + * search), and it takes effect from the next search rather than reshuffling what + * is already on screen. The model is not a choice - there is one reranker - so + * the only field is which key pays, and "off" is one of its options. + */ +export function RerankDialog({ + open, + onOpenChange, + rerankSecretId, + collectionName, + onSave, +}: RerankDialogProps) { + const tErrors = useTranslations("errors"); + const t = useTranslations("kb"); + const { secrets } = useSecrets(); + const rerankKeys = secrets.filter((secret) => secret.purpose === RERANK_KEY_PURPOSE); + const [draftSecretId, setDraftSecretId] = useState(rerankSecretId); + const [isSaving, setIsSaving] = useState(false); + const [errors, setErrors] = useState>>({}); + + // Reopening shows what the server holds, not a draft abandoned last time - and + // a change made elsewhere must not leave a stale pick to be posted back over. + const opened = useChanged(open); + const secretMoved = useChanged(rerankSecretId); + if (opened || secretMoved) { + if (open) { + setDraftSecretId(rerankSecretId); + setErrors({}); + } + } + + const changed = draftSecretId !== rerankSecretId; + + const handleSave = async () => { + setIsSaving(true); + try { + // The pair the backend reads together: a key turns reranking on with the + // one model there is, no key turns it off. Both are always sent, so the + // update is unambiguously "change reranking" rather than "leave it". + await onSave({ + rerank_model: draftSecretId ? DEFAULT_RERANK_MODEL : null, + rerank_secret_id: draftSecretId, + }); + onOpenChange(false); + } catch (err) { + const failure = submitFailure(err, { fields: ["rerank_secret_id"] }, tErrors); + setErrors(failure.fields); + if (failure.toast) toast.error(failure.toast); + } finally { + setIsSaving(false); + } + }; + + return ( + + + + {t("rerankSettings")} + + {t.rich("rerankSettingsDescription", { + name: collectionName, + mono: (chunks) => {chunks}, + })} + + + +
+ + + {errors.rerank_secret_id && ( +

{errors.rerank_secret_id}

+ )} +

{t("rerankHelp")}

+ +
+ + + + + + + + ); +} diff --git a/frontend/src/components/kb/rerank-panel.test.tsx b/frontend/src/components/kb/rerank-panel.test.tsx new file mode 100644 index 000000000..9ad88e723 --- /dev/null +++ b/frontend/src/components/kb/rerank-panel.test.tsx @@ -0,0 +1,106 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen } from "@testing-library/react"; +import type { ReactNode } from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { RerankPanel } from "./rerank-panel"; +import { apiClient } from "@/lib/api-client"; +import { DEFAULT_INGESTION_CONFIG } from "@/lib/ingestion-config"; +import type { KnowledgeBase } from "@/types/knowledge-base"; + +/** + * The panel is what tells someone reading a collection whether its searches are + * reranked, and with which key - a retrieval-time fact that changes on a + * different day than how the documents were read, which is why it is its own + * section rather than a line in the ingestion panel. + */ + +vi.mock("@/lib/api-client", async (importOriginal) => ({ + ...(await importOriginal()), + apiClient: { get: vi.fn(), post: vi.fn(), patch: vi.fn(), delete: vi.fn() }, +})); + +const SECRETS = { + items: [{ id: "co-1", name: "Cohere prod", hint: "4242", purpose: "cohere", kind: "api_key" }], + total: 1, +}; + +function kb(overrides: Partial = {}): KnowledgeBase { + return { + id: "kb-1", + organization_id: "org-1", + owner_user_id: null, + name: "Handbook", + description: null, + scope: "org", + collection_name: "handbook_a1b2c3", + is_default: false, + ingestion_config: DEFAULT_INGESTION_CONFIG, + embedding_model: "text-embedding-3-large", + embedding_dim: 3072, + embedding_provider: "openai", + embedding_secret_id: null, + rerank_model: null, + rerank_secret_id: null, + created_at: "2026-07-01T00:00:00Z", + updated_at: null, + document_count: 0, + indexed_count: 0, + chunk_count: 0, + ...overrides, + }; +} + +function wrapper({ children }: { children: ReactNode }) { + const client = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + return {children}; +} + +beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(apiClient.get).mockImplementation(async (path: string) => { + if (path === "/secrets") return SECRETS; + return { items: [], total: 0 }; + }); +}); + +describe("what the panel says", () => { + it("says reranking is off, and that distance alone orders the results", () => { + render(, { wrapper }); + expect(screen.getByText(/ordered by vector distance/)).toBeInTheDocument(); + }); + + it("names the model and the key when it is on", async () => { + render(, { + wrapper, + }); + expect(screen.getByText("rerank-v3.5")).toBeInTheDocument(); + expect(await screen.findByText(/billed to Cohere prod/)).toBeInTheDocument(); + }); + + it("falls back to a neutral key label when the reader cannot list secrets", () => { + // A `collections:edit` holder need not hold `connections:manage`, so + // `GET /secrets` answers 403 and an empty list - the key's id resolves to no + // name, and the panel must still say reranking is on rather than break. + vi.mocked(apiClient.get).mockResolvedValue({ items: [], total: 0 }); + render(, { + wrapper, + }); + expect(screen.getByText(/billed to a Cohere key/)).toBeInTheDocument(); + }); +}); + +describe("the edit affordance", () => { + it("offers Edit to a caller who may write", async () => { + const onEdit = vi.fn(); + render(, { wrapper }); + const { default: userEvent } = await import("@testing-library/user-event"); + await userEvent.click(screen.getByRole("button", { name: "Edit" })); + expect(onEdit).toHaveBeenCalledOnce(); + }); + + it("shows no Edit to a caller who may not", () => { + render(, { wrapper }); + expect(screen.queryByRole("button", { name: "Edit" })).toBeNull(); + }); +}); diff --git a/frontend/src/components/kb/rerank-panel.tsx b/frontend/src/components/kb/rerank-panel.tsx new file mode 100644 index 000000000..aaa43b17d --- /dev/null +++ b/frontend/src/components/kb/rerank-panel.tsx @@ -0,0 +1,72 @@ +"use client"; + +import { SlidersHorizontal } from "lucide-react"; + +import { Button } from "@/components/ui"; +import { useSecrets } from "@/hooks"; +import type { KnowledgeBase } from "@/types"; +import { useTranslations } from "next-intl"; + +interface RerankPanelProps { + kb: KnowledgeBase; + /** Absent when the caller may not write - the panel is then facts only. */ + onEdit?: () => void; +} + +/** + * Whether this collection reranks its search results, and with what. + * + * A retrieval-time setting, so it is its own section rather than a line in "how + * documents are read": that panel describes what was baked into the collection + * at ingestion, this describes what happens to a query against it now, and the + * two change on different days for different reasons. + * + * Reranking is configured (both a model and a key) or it is off; there is no + * half state, because the backend resolves anything but a usable key to no + * reranker. The key's name is resolved from the vault, and falls back to a + * neutral label when the reader cannot list secrets (a `collections:edit` + * holder need not hold `connections:manage`). + */ +export function RerankPanel({ kb, onEdit }: RerankPanelProps) { + const t = useTranslations("kb"); + const { secrets } = useSecrets(); + const configured = Boolean(kb.rerank_model && kb.rerank_secret_id); + const keyName = secrets.find((secret) => secret.id === kb.rerank_secret_id)?.name; + + return ( +
+
+

+ {t("rerank")} +

+ {onEdit && ( + + )} +
+ +
+ {configured ? ( +
+ + {kb.rerank_model} + + {t("rerankBilledTo", { key: keyName ?? t("rerankKeyConfigured") })} + + +

+ {t("rerankOnExplained")} +

+
+ ) : ( +

{t("rerankOffExplained")}

+ )} +
+
+ ); +} diff --git a/frontend/src/components/kb/reusable-integrations.integration.test.tsx b/frontend/src/components/kb/reusable-integrations.integration.test.tsx index 5feef8921..184cc6a51 100644 --- a/frontend/src/components/kb/reusable-integrations.integration.test.tsx +++ b/frontend/src/components/kb/reusable-integrations.integration.test.tsx @@ -58,6 +58,8 @@ function kb(id: string, name: string, collection: string): KnowledgeBase { embedding_provider: "openrouter", embedding_secret_id: null, embedding_dim: 3072, + rerank_model: null, + rerank_secret_id: null, created_at: "2026-07-01T00:00:00Z", updated_at: null, document_count: 0, diff --git a/frontend/src/hooks/use-knowledge-bases.test.tsx b/frontend/src/hooks/use-knowledge-bases.test.tsx index a9e8e9a9d..8952123d3 100644 --- a/frontend/src/hooks/use-knowledge-bases.test.tsx +++ b/frontend/src/hooks/use-knowledge-bases.test.tsx @@ -585,6 +585,48 @@ describe("one collection's page", () => { ); }); + it("saves the rerank pair together and keeps the collection it was handed back", async () => { + serveDetail(); + const { result } = renderHook(() => useKBDetail("kb-1"), { wrapper }); + await waitFor(() => expect(result.current.kb).toMatchObject({ id: "kb-1" })); + vi.mocked(apiClient.patch).mockResolvedValue({ + id: "kb-1", + name: "Handbook", + rerank_model: "rerank-v3.5", + }); + + await act(async () => { + await result.current.updateRerank({ + rerank_model: "rerank-v3.5", + rerank_secret_id: "co-1", + }); + }); + + expect(apiClient.patch).toHaveBeenCalledWith("/kb/kb-1", { + rerank_model: "rerank-v3.5", + rerank_secret_id: "co-1", + }); + await waitFor(() => expect(result.current.kb).toMatchObject({ rerank_model: "rerank-v3.5" })); + expect(toast.success).toHaveBeenCalledWith("Reranking updated"); + }); + + it("refuses to change reranking with no collection open", async () => { + const { result } = renderHook(() => useKBDetail(null), { wrapper }); + + await expect( + result.current.updateRerank({ rerank_model: null, rerank_secret_id: null }), + ).rejects.toThrow("No knowledge base is open"); + }); + + it("lets a refused rerank key through to the dialog that owns the picker", async () => { + const { result } = renderHook(() => useKBDetail("kb-1"), { wrapper }); + vi.mocked(apiClient.patch).mockRejectedValue(new Error("That key is for tavily")); + + await expect( + result.current.updateRerank({ rerank_model: "rerank-v3.5", rerank_secret_id: "co-1" }), + ).rejects.toThrow("That key is for tavily"); + }); + it("drops a deleted document and the count with it", async () => { serveDetail({ documents: [document("d-1"), document("d-2")], documentsTotal: 2 }); const { result } = renderHook(() => useKBDetail("kb-1"), { wrapper }); diff --git a/frontend/src/hooks/use-knowledge-bases.ts b/frontend/src/hooks/use-knowledge-bases.ts index 68fc15c26..3db369a5d 100644 --- a/frontend/src/hooks/use-knowledge-bases.ts +++ b/frontend/src/hooks/use-knowledge-bases.ts @@ -29,6 +29,7 @@ import type { KBDocumentList, KnowledgeBase, KnowledgeBaseList, + UpdateRerankInput, } from "@/types"; export function useKnowledgeBases() { @@ -334,6 +335,30 @@ export function useKBDetail(id: string | null) { [id, activeOrgId, stillSameTenant, queryClient, t], ); + /** + * Turn reranking on, change its key, or turn it off - from now on. + * + * The pair is sent together, which is how the backend tells "change reranking" + * from "leave it alone": an update carrying both fields sets them, one + * carrying neither does not. `null`/`null` is the off signal. Nothing already + * retrieved changes; this governs the next search. The refusal is rethrown so + * the dialog can put "that key is for something else" beside the picker. + */ + const updateRerank = useCallback( + async (input: UpdateRerankInput): Promise => { + // i18n-exempt: a narrowing guard on a hook only mounted with an id, not copy. + if (!id) throw new Error("No knowledge base is open"); + const startedIn = activeOrgId; + const updated = await apiClient.patch(`/kb/${id}`, input); + if (stillSameTenant(startedIn)) { + queryClient.setQueryData(qk.kb.detail(id), updated); + toast.success(t("rerankSaved")); + } + return updated; + }, + [id, activeOrgId, stillSameTenant, queryClient, t], + ); + /** * Move this collection's embeddings to another provider, or another key. * @@ -620,6 +645,7 @@ export function useKBDetail(id: string | null) { refresh, loadMoreDocuments, updateIngestion, + updateRerank, updateEmbeddings, uploadDocument, deleteDocument, diff --git a/frontend/src/hooks/use-reusable-integrations.test.tsx b/frontend/src/hooks/use-reusable-integrations.test.tsx index e9e859b81..4c8600a37 100644 --- a/frontend/src/hooks/use-reusable-integrations.test.tsx +++ b/frontend/src/hooks/use-reusable-integrations.test.tsx @@ -57,6 +57,8 @@ const TARGET: KnowledgeBase = { embedding_provider: "openrouter", embedding_secret_id: null, embedding_dim: 3072, + rerank_model: null, + rerank_secret_id: null, created_at: "2026-07-01T00:00:00Z", updated_at: null, document_count: 0, diff --git a/frontend/src/lib/onboarding/flows.test.ts b/frontend/src/lib/onboarding/flows.test.ts index fb4895c37..eb8edf09e 100644 --- a/frontend/src/lib/onboarding/flows.test.ts +++ b/frontend/src/lib/onboarding/flows.test.ts @@ -105,6 +105,7 @@ describe("FLOWS", () => { "flow-kb-field-name", "flow-kb-field-scope", "flow-kb-field-embeddings", + "flow-kb-field-rerank", "flow-kb-field-create", "flow-agent-knowledge-return-nav", "flow-agent-knowledge-return-edit", @@ -402,6 +403,7 @@ describe("stepsForFlow", () => { "flow-kb-field-name", "flow-kb-field-scope", "flow-kb-field-embeddings", + "flow-kb-field-rerank", "flow-kb-field-create", "flow-agent-knowledge-return-nav", "flow-agent-knowledge-return-edit", diff --git a/frontend/src/lib/onboarding/flows.ts b/frontend/src/lib/onboarding/flows.ts index 1ad335a65..cd8516d3c 100644 --- a/frontend/src/lib/onboarding/flows.ts +++ b/frontend/src/lib/onboarding/flows.ts @@ -341,6 +341,15 @@ function kbDialogSteps(page: string, requires?: string): FlowStep[] { blockSubmit: "kb-dialog-create", requires, }, + { + id: "flow-kb-field-rerank", + page, + target: "kb-dialog-rerank", + permission: Perm.collectionsEdit, + inOverlay: true, + blockSubmit: "kb-dialog-create", + requires, + }, { id: "flow-kb-field-create", page, diff --git a/frontend/src/lib/onboarding/tour.test.ts b/frontend/src/lib/onboarding/tour.test.ts index 940c2d292..04d84f244 100644 --- a/frontend/src/lib/onboarding/tour.test.ts +++ b/frontend/src/lib/onboarding/tour.test.ts @@ -45,7 +45,7 @@ const BUILDER_STEPS = [ // Every collection-detail stop, in the order the "?" walks them. The launch // pass takes one of these (kb-documents). -const KB_STEPS = ["kb-header", "kb-documents", "kb-ingestion", "kb-sync"]; +const KB_STEPS = ["kb-header", "kb-documents", "kb-ingestion", "kb-rerank", "kb-sync"]; // The organization detail walk: the members page (profile, then the list), then // across into the roles matrix. None are inTour — orgs is a "?"-only section. diff --git a/frontend/src/lib/onboarding/tour.ts b/frontend/src/lib/onboarding/tour.ts index d64b4f6ae..ff5d9d1a9 100644 --- a/frontend/src/lib/onboarding/tour.ts +++ b/frontend/src/lib/onboarding/tour.ts @@ -340,6 +340,13 @@ export const TOUR_STEPS: readonly TourStep[] = [ activate: "kb-tab-ingestion", permission: Perm.collectionsView, }, + { + id: "kb-rerank", + page: KB_DETAIL, + target: "kb-rerank", + activate: "kb-tab-ingestion", + permission: Perm.collectionsView, + }, { id: "kb-sync", page: KB_DETAIL, diff --git a/frontend/src/lib/rerank-config.ts b/frontend/src/lib/rerank-config.ts new file mode 100644 index 000000000..d324fd2af --- /dev/null +++ b/frontend/src/lib/rerank-config.ts @@ -0,0 +1,20 @@ +/** + * The three constants reranking is configured with, in one place because two + * surfaces set it: the create dialog and the detail page's edit dialog. Keeping + * them apart is how the two drift - one dialog offering a model the other does + * not, a purpose filter that stops matching the key it stored. + * + * All three mirror the backend. `RERANK_KEY_PURPOSE` is the single entry in + * `RERANK_KEY_PURPOSES` (`app/services/rerank_resolution.py`); `DEFAULT_RERANK_MODEL` + * is the value stored as `rerank_model`. No endpoint lists rerankers - there is + * one - so the model is a constant rather than a fetched list. + */ + +/** The purpose a vault key must carry to pay for reranking. */ +export const RERANK_KEY_PURPOSE = "cohere"; + +/** Sentinel for "no reranking" - a Select item may not have an empty value. */ +export const RERANK_OFF = "__off__"; + +/** The one reranker there is, stored as the collection's `rerank_model`. */ +export const DEFAULT_RERANK_MODEL = "rerank-v3.5"; diff --git a/frontend/src/types/knowledge-base.ts b/frontend/src/types/knowledge-base.ts index 5c0f6f78d..59a4e97e3 100644 --- a/frontend/src/types/knowledge-base.ts +++ b/frontend/src/types/knowledge-base.ts @@ -108,6 +108,14 @@ export interface KnowledgeBase { */ embedding_model: string; embedding_dim: number; + /** + * The reranker reordering this collection's search results, and the org vault + * key it is billed to. Both null unless reranking is configured. Unlike the + * embedding model this pair can be changed after creation - see + * `UpdateRerankInput`. + */ + rerank_model: string | null; + rerank_secret_id: string | null; /** * Whose endpoint serves that model - a provider id from * `GET /rag/embedding-models`. @@ -161,6 +169,27 @@ export interface CreateKnowledgeBaseInput { embedding_provider?: string; /** The org vault key that pays for embeddings; omit for the deployment key. */ embedding_secret_id?: string; + /** + * The reranker applied to search results. Reranking is on only when this and + * `rerank_secret_id` are both sent; omit both to leave it off. + */ + rerank_model?: string; + /** The org vault key that pays for reranking - a `cohere`-purpose api_key. */ + rerank_secret_id?: string; +} + +/** + * Changing a collection's reranking after creation. + * + * The pair is read together on the backend: send both to turn reranking on or + * change its key, both `null` to turn it off, and omit both to leave it be + * (which is why they are `null`-able rather than merely optional - `null` is the + * "off" signal, absence is "don't touch"). No other field changes here; name, + * description and ingestion have their own paths. + */ +export interface UpdateRerankInput { + rerank_model: string | null; + rerank_secret_id: string | null; } /** What a collection's embeddings may be re-pointed at after the fact. */ diff --git a/frontend/src/types/secrets.ts b/frontend/src/types/secrets.ts index 3e282710c..22074c9f7 100644 --- a/frontend/src/types/secrets.ts +++ b/frontend/src/types/secrets.ts @@ -97,7 +97,7 @@ export interface SecretKindList { /** A stored secret, identified by everything except what it holds. */ /** One place a secret is bound, so "can I delete this" has an answer. */ export interface SecretUsage { - kind: "agent"; + kind: "agent" | "knowledge_base"; id: string; name: string; } diff --git a/frontend/src/types/stats.ts b/frontend/src/types/stats.ts index 2ac0277cd..e41c5e753 100644 --- a/frontend/src/types/stats.ts +++ b/frontend/src/types/stats.ts @@ -61,14 +61,17 @@ export interface ProviderCost { export interface CostBlock { /** * Serialised Decimals. `period_usd` is the whole bill - models plus - * ingestion - and the two halves below sum to it; the calendar - * month-to-date figure still lives on GET /spend. + * ingestion plus retrieval - and the three parts below sum to it; the + * calendar month-to-date figure still lives on GET /spend. */ period_usd: string; previous_period_usd: string; model_usd: string; /** Zero at scope=own: a document is indexed by a worker, for nobody. */ ingestion_usd: string; + /** What a metered /rag/search spent on embeddings and reranking, kept apart + * from indexing so a search is not reported as it. Zero at scope=own. */ + retrieval_usd: string; by_provider: ProviderCost[]; }