From 1292250f130a8a90ac00f1c060265c53c67586d1 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Tue, 18 Aug 2026 10:24:50 +0200 Subject: [PATCH 01/37] feat(budget): book a pre-priced entry to the active ledger record_ambient_usage() prices every call through genai-prices, which knows chat and embedding models and nothing else. A reranker call routed through it would book cost_usd=0, priced=False, so its spend would be invisible to the monthly budgets and reported as a floor. book_ambient_spend(entry) is the sibling for spend that is not token-priced: the caller computes the cost from a published per-search price and hands the finished SpendEntry over, so it lands priced=True with a real number. Like its sibling it is a no-op when nothing is metering - a search outside any run has no ledger open and must still run rather than refuse. Groundwork for the RAG reranker (#142); no caller yet. Refs #142 --- .../agents/capabilities/budget/__init__.py | 2 + .../agents/capabilities/budget/_capability.py | 18 ++++++++ backend/tests/test_spend.py | 42 +++++++++++++++++++ 3 files changed, 62 insertions(+) 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/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.""" From 6e02b1d267f0d451c27f032f7466027e5c662985 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Tue, 18 Aug 2026 10:31:00 +0200 Subject: [PATCH 02/37] feat(rag): add rerank columns to knowledge_bases Two nullable columns mirroring the embedding pair: rerank_model (the reranker's name) and rerank_secret_id (the org vault key that pays for it, FK to organization_secrets, SET NULL on delete). Reranking is on for a collection only when both are set; either NULL leaves retrieval exactly as it was, so existing rows and unconfigured deployments are unchanged. Unlike the embedding key there is no deployment fallback - a reranker with no key is simply off - so nothing goes into RAGSettings. Verified against a real pgvector database: alembic upgrade head, alembic check reports no drift (model matches the migration), and downgrade -1 -> upgrade round-trips. tests/test_migrations.py green. Refs #142 --- .../versions/0037_knowledge_base_rerank.py | 61 +++++++++++++++++++ backend/app/db/models/knowledge_base.py | 13 ++++ 2 files changed, 74 insertions(+) create mode 100644 backend/alembic/versions/0037_knowledge_base_rerank.py diff --git a/backend/alembic/versions/0037_knowledge_base_rerank.py b/backend/alembic/versions/0037_knowledge_base_rerank.py new file mode 100644 index 000000000..58cf982db --- /dev/null +++ b/backend/alembic/versions/0037_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: 0037_knowledge_base_rerank +Revises: 0036_conversation_reminder_state +Create Date: 2026-08-18 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa + +from alembic import op + +revision: str = "0037_knowledge_base_rerank" +down_revision: str | None = "0036_conversation_reminder_state" +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/app/db/models/knowledge_base.py b/backend/app/db/models/knowledge_base.py index f10a4da8a..b60f50c81 100644 --- a/backend/app/db/models/knowledge_base.py +++ b/backend/app/db/models/knowledge_base.py @@ -58,6 +58,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( From 63cf6f610ec630761c6327ccc0c6786af5019602 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Tue, 18 Aug 2026 10:40:35 +0200 Subject: [PATCH 03/37] feat(rag): resolve a collection's reranker and whose key pays reranker_for_collection() is the sibling of embeddings_for_collection(): it asks per collection whether a reranker is configured and returns its model and the organization key it runs on, or None. The one deliberate difference from embedding resolution is the whole design of the feature being off by default. Embeddings fall back to the deployment key when a collection's chosen one is gone; reranking has no deployment key, so every path but a usable organization secret resolves to None and retrieval is byte-for-byte its pre-feature self. The three degraded reasons (secret missing / unusable / wrong kind) are still told apart from the normal off state and logged, exactly as EmbeddingKeySource names its own - a chosen key that vanished is an operator's problem, a collection that chose nothing is not. Adds the `cohere` secret purpose (category other, api_key) with RERANK_KEY_PURPOSES as its consumer. rerank_resolution.py joins the coverage + ty gates beside embedding_resolution.py (100%, verified). Refs #142 --- backend/app/core/catalog/services.json | 8 ++ backend/app/services/rerank_resolution.py | 144 ++++++++++++++++++++ backend/pyproject.toml | 2 + backend/tests/test_coverage_gate.py | 1 + backend/tests/test_rerank_resolution.py | 156 ++++++++++++++++++++++ 5 files changed, 311 insertions(+) create mode 100644 backend/app/services/rerank_resolution.py create mode 100644 backend/tests/test_rerank_resolution.py diff --git a/backend/app/core/catalog/services.json b/backend/app/core/catalog/services.json index b71604f25..781845520 100644 --- a/backend/app/core/catalog/services.json +++ b/backend/app/core/catalog/services.json @@ -53,5 +53,13 @@ "kind": "api_key", "help_url": "https://app.daytona.io/dashboard/keys", "description": "Cloud sandboxes an agent works in, billed to this organization's own Daytona account." + }, + { + "id": "cohere", + "label": "Cohere", + "category": "other", + "kind": "api_key", + "help_url": "https://dashboard.cohere.com/api-keys", + "description": "Rerank knowledge-search results with Cohere, billed to this organization's own key." } ] diff --git a/backend/app/services/rerank_resolution.py b/backend/app/services/rerank_resolution.py new file mode 100644 index 000000000..74d72345d --- /dev/null +++ b/backend/app/services/rerank_resolution.py @@ -0,0 +1,144 @@ +"""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 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; +# the tuple exists so a second provider is one entry, not a hunt, and so the +# `cohere` entry in the services catalog has a consumer that names it. +RERANK_KEY_PURPOSES = ("cohere",) + + +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) -> 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. + """ + async with get_db_context() as db: + kb = await knowledge_base_repo.get_by_collection_name(db, collection_name) + 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 or secret_id is None or organization_id is None: + return None, RerankKeySource.NOT_CONFIGURED + + 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/pyproject.toml b/backend/pyproject.toml index 1dad236e4..282a26adf 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -519,6 +519,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", @@ -700,6 +701,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 diff --git a/backend/tests/test_coverage_gate.py b/backend/tests/test_coverage_gate.py index 9debb1f97..a610cec1f 100644 --- a/backend/tests/test_coverage_gate.py +++ b/backend/tests/test_coverage_gate.py @@ -86,6 +86,7 @@ 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/mcp_catalog.py", diff --git a/backend/tests/test_rerank_resolution.py b/backend/tests/test_rerank_resolution.py new file mode 100644 index 000000000..02c94f0b6 --- /dev/null +++ b/backend/tests/test_rerank_resolution.py @@ -0,0 +1,156 @@ +"""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_by_collection_name = AsyncMock(return_value=kb) + secrets.get = AsyncMock(return_value=secret_row) + return await reranker_for_collection("handbook"), 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_is_off(self): + resolved, secrets = await _resolve(_kb(model="rerank-v3.5", secret_id=None)) + assert resolved is None + secrets.get.assert_not_called() + + 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) + + +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, + } From 5258ec26b59a6d054c246eb18e45fbba6fd160e3 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Tue, 18 Aug 2026 10:45:37 +0200 Subject: [PATCH 04/37] feat(rag): add BaseReranker and a Cohere implementation BaseReranker is the interface retrieval depends on; CohereReranker is the first and only implementation. rerank() reorders the candidate SearchResults by Cohere's relevance score and re-scores them with it, so a caller ordering or thresholding on score reads the reranker's judgement rather than the vector distance the candidates arrived with. Cost is booked through book_ambient_spend, not record_ambient_usage: a rerank call is priced per search unit (one query, up to 100 documents), not per token, and genai-prices does not know rerank models. The per-search price is the one number in the metering path that lives in this repository - $0.002/unit, confirmed against cohere.com/pricing on 2026-08-18, with a comment telling the next maintainer to re-check the page and the constant together. Booked only after the call returns, because Cohere does not bill a failed request; a raise propagates so retrieval can degrade to the un-reranked order rather than failing the search. The client is built lazily and injectable, so the tests drive it with no network and no key. Refs #142 --- backend/app/services/rag/reranker.py | 122 +++++++++++++++++++++++++++ backend/tests/test_reranker.py | 110 ++++++++++++++++++++++++ 2 files changed, 232 insertions(+) create mode 100644 backend/app/services/rag/reranker.py create mode 100644 backend/tests/test_reranker.py diff --git a/backend/app/services/rag/reranker.py b/backend/app/services/rag/reranker.py new file mode 100644 index 000000000..ebf4c16ab --- /dev/null +++ b/backend/app/services/rag/reranker.py @@ -0,0 +1,122 @@ +"""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 + +import cohere + +from app.agents.capabilities.budget import SpendEntry, book_ambient_spend +from app.services.rag.models import SearchResult + +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. A request with more is split and billed as several - a query with +# 250 documents is three units. +_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. + """ + + def __init__(self, model: str, api_key: str, client: AsyncClientV2 | None = None) -> None: + self.model = model + self._api_key = api_key + self._client = client + + @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 [] + + response = await self.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(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 + ] + + def _spend_entry(self, document_count: int) -> SpendEntry: + units = ceil(document_count / _DOCS_PER_SEARCH_UNIT) + return SpendEntry( + model_name=self.model, + input_tokens=0, + output_tokens=0, + cost_usd=_PRICE_PER_SEARCH_UNIT_USD * units, + priced=True, + ) diff --git a/backend/tests/test_reranker.py b/backend/tests/test_reranker.py new file mode 100644 index 000000000..46e15146f --- /dev/null +++ b/backend/tests/test_reranker.py @@ -0,0 +1,110 @@ +"""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 + +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 + +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]) -> AsyncMock: + """A fake Cohere client whose rerank returns these (index, score) items.""" + response = SimpleNamespace( + results=[SimpleNamespace(index=index, relevance_score=score) for index, score in ranked] + ) + 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): + 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_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 From d88b8aa053db08afbdcb7b8e1798d4d4ae8c86f1 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Tue, 18 Aug 2026 10:52:17 +0200 Subject: [PATCH 05/37] feat(rag): rerank retrieval candidates when a collection is configured RetrievalService takes an optional reranker resolver. When a collection resolves one, retrieve() overfetches a wider candidate net (4x rather than 2x), reranks, and truncates to the limit; retrieve_multi() gathers each collection's candidates and reranks the union once, because an agent's bound collections share one organization and so one reranker. With no resolver - or none configured for the collection - every path is byte-for-byte the previous by-distance one, down to each collection contributing its top `limit` before a multi-collection merge. Recall is split out of retrieve() into _recall() so the multi path can fuse before ranking rather than rank-then-fuse. The collection stamp moves onto every candidate so it survives reranking, which builds fresh results. A reranker failure at query time degrades to the distance order with a log line rather than failing the search - reranking is an improvement on a working retrieval, not a dependency of it; the misconfiguration cases never reach here, resolution having already turned those into no reranker. deps wires reranker_for_collection through a composition-root adapter that binds the resolved credential to a CohereReranker - the one place a second provider would branch. Refs #142 --- backend/app/api/deps.py | 22 ++++- backend/app/services/rag/retrieval.py | 114 ++++++++++++++++++---- backend/tests/test_retrieval_reranking.py | 98 +++++++++++++++++++ 3 files changed, 216 insertions(+), 18 deletions(-) create mode 100644 backend/tests/test_retrieval_reranking.py diff --git a/backend/app/api/deps.py b/backend/app/api/deps.py index f6913b433..31cabbc7b 100644 --- a/backend/app/api/deps.py +++ b/backend/app/api/deps.py @@ -857,9 +857,11 @@ async def verify_api_key( from app.services.embedding_resolution import embeddings_for_collection from app.services.rag.ingestion import IngestionService from app.services.rag.documents import DocumentProcessor +from app.services.rag.reranker import BaseReranker, CohereReranker from app.services.rag.retrieval import RetrievalService from app.services.rag.vectorstore import PgVectorStore from app.services.rag.vectorstore import BaseVectorStore +from app.services.rerank_resolution import reranker_for_collection def get_embedding_service(request: Request) -> EmbeddingService: @@ -884,9 +886,27 @@ def get_vectorstore(request: Request, embedder: EmbeddingSvc) -> BaseVectorStore VectorStoreSvc = Annotated[BaseVectorStore, Depends(get_vectorstore)] +async def _reranker_for_collection(collection_name: str) -> BaseReranker | None: + """Bind a collection's resolved reranker credential to a concrete reranker. + + The composition root for reranking: resolution answers whether a collection + is configured and with whose key, and this turns that into the one + implementation there is. A second provider is a branch here, not a change to + retrieval. + """ + resolved = await reranker_for_collection(collection_name) + if resolved is None: + return None + return CohereReranker(model=resolved.model, api_key=resolved.api_key) + + def get_retrieval_service(vector_store: VectorStoreSvc) -> RetrievalService: """Create RetrievalService instance.""" - return RetrievalService(vector_store=vector_store, settings=settings.rag) + return RetrievalService( + vector_store=vector_store, + settings=settings.rag, + reranker_resolver=_reranker_for_collection, + ) RetrievalSvc = Annotated[RetrievalService, Depends(get_retrieval_service)] diff --git a/backend/app/services/rag/retrieval.py b/backend/app/services/rag/retrieval.py index 1ff0685c0..83726fb2d 100644 --- a/backend/app/services/rag/retrieval.py +++ b/backend/app/services/rag/retrieval.py @@ -4,15 +4,29 @@ import logging import time from abc import ABC, abstractmethod +from collections.abc import Awaitable, Callable 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], 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: @@ -38,10 +52,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( @@ -105,6 +123,40 @@ async def _bm25_search( if s > 0 ] + async def _reranker_for(self, collection_name: str) -> 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. + """ + if self._reranker_resolver is None: + return None + return await self._reranker_resolver(collection_name) + + @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, @@ -113,9 +165,31 @@ async def retrieve( min_score: float = 0.0, filter: str = "", ) -> list[SearchResult]: - # Overfetch so min-score filtering and dedup still leave `limit` results. - fetch_multiplier = 2 + reranker = await self._reranker_for(collection_name) + multiplier = _RERANK_FETCH_MULTIPLIER if reranker else _DEFAULT_FETCH_MULTIPLIER + candidates = await self._recall( + query, collection_name, limit, min_score, filter, fetch_multiplier=multiplier + ) + 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, + ) -> 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. + """ logger.info( "[RETRIEVAL] Query: '%.50s...', collection: %s, limit: %d, filter: '%s'", query, @@ -179,23 +253,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, @@ -214,17 +288,23 @@ 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: the bound collections of + one agent share one organization and so one reranker, and reranking each + collection separately then merging the winners would rank against the + wrong pool. Absent a reranker this is byte-for-byte the previous merge - + each collection's top `limit`, fused, sorted, deduplicated, truncated. """ + reranker = await self._reranker_for(collection_names[0]) if collection_names else None + 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, - ) + recalled = await self._recall( + query, name, limit, min_score, "", fetch_multiplier=multiplier ) + all_results.extend(recalled if reranker else recalled[:limit]) all_results.sort(key=lambda r: r.score, reverse=True) @@ -236,4 +316,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/tests/test_retrieval_reranking.py b/backend/tests/test_retrieval_reranking.py new file mode 100644 index 000000000..08c01dafc --- /dev/null +++ b/backend/tests/test_retrieval_reranking.py @@ -0,0 +1,98 @@ +"""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 + +import pytest + +from app.services.rag.models import SearchResult +from app.services.rag.reranker import BaseReranker +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 _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) + 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) + 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) + 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) + 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 + ) + 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) + assert results[0].metadata["collection"] == "handbook" From 8dd6ca8f406ae8d7b5fe5c4b9e5cc01168b2e6d1 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Tue, 18 Aug 2026 11:02:28 +0200 Subject: [PATCH 06/37] feat(rag): meter POST /rag/search against the organization The search route embedded the query - and, once a collection is configured, reranks - inside no metered_by block and against no ledger, so neither cost reached the organization's monthly bill (#16 class). Reranking is what made that worth fixing; metering the embeddings too is the beneficial side effect. KnowledgeSearchService owns it: it resolves collection access, opens a ledger scoped to the caller's organization, runs the search inside metered_by so the ambient embedding and rerank calls book to it, and persists what they spent to ingestion_spend with a null document id - the same sink a worker's ingestion spend lands in. The route drops to plumbing, which also keeps the metering out of a route handler the layering forbids logic in. Access is resolved before the ledger opens, so a cross-tenant collection still refuses the whole search before any vectors are read - the tenant-isolation route tests in test_platform_flows.py stay green. knowledge_search.py joins the coverage + ty gates (100%). Refs #142, #16 --- backend/app/api/deps.py | 11 ++ backend/app/api/routes/v1/rag.py | 26 +---- backend/app/services/knowledge_search.py | 103 ++++++++++++++++++ backend/pyproject.toml | 2 + backend/tests/test_coverage_gate.py | 1 + backend/tests/test_knowledge_search.py | 127 +++++++++++++++++++++++ 6 files changed, 249 insertions(+), 21 deletions(-) create mode 100644 backend/app/services/knowledge_search.py create mode 100644 backend/tests/test_knowledge_search.py diff --git a/backend/app/api/deps.py b/backend/app/api/deps.py index 31cabbc7b..3ce995f94 100644 --- a/backend/app/api/deps.py +++ b/backend/app/api/deps.py @@ -857,6 +857,7 @@ async def verify_api_key( from app.services.embedding_resolution import embeddings_for_collection 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 BaseReranker, CohereReranker from app.services.rag.retrieval import RetrievalService from app.services.rag.vectorstore import PgVectorStore @@ -912,6 +913,16 @@ def get_retrieval_service(vector_store: VectorStoreSvc) -> RetrievalService: 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) diff --git a/backend/app/api/routes/v1/rag.py b/backend/app/api/routes/v1/rag.py index 3c9fdb4af..d513305cf 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, @@ -242,33 +242,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/services/knowledge_search.py b/backend/app/services/knowledge_search.py new file mode 100644 index 000000000..b648d4717 --- /dev/null +++ b/backend/app/services/knowledge_search.py @@ -0,0 +1,103 @@ +"""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.repositories import ingestion_spend_repo + +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. + """ + names = request.collection_names or [request.collection_name] + collections = [kb.collection_name for kb in await self.access.readable_all(ctx, names)] + + ledger = SpendLedger(organization_id=ctx.organization_id) + 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, + ) + 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 "", + ) + + await self._record_spend(ledger) + return results + + async def _record_spend(self, 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. + """ + 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( + self.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), + ) diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 282a26adf..932e8ad78 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -531,6 +531,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/mcp_connection.py", "app/services/model_profile.py", @@ -734,6 +735,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/mcp_connection.py", "app/services/model_profile.py", diff --git a/backend/tests/test_coverage_gate.py b/backend/tests/test_coverage_gate.py index a610cec1f..f1a133765 100644 --- a/backend/tests/test_coverage_gate.py +++ b/backend/tests/test_coverage_gate.py @@ -89,6 +89,7 @@ def _matches_glob(path: str, pattern: str) -> bool: "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_knowledge_search.py b/backend/tests/test_knowledge_search.py new file mode 100644 index 000000000..a276ad31c --- /dev/null +++ b/backend/tests/test_knowledge_search.py @@ -0,0 +1,127 @@ +"""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 SpendEntry, book_ambient_spend +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" + + +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 + + 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"} From 22a2395dd4569d60771d468d57fe5e3eef46c65d Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Tue, 18 Aug 2026 11:12:54 +0200 Subject: [PATCH 07/37] feat(rag): let a knowledge base set its reranker The rerank model and key are now on KnowledgeBaseCreate/Update/Read and threaded through the service and repository, so a collection can actually be configured to rerank. Two rules the service enforces: - A model and a key together, or neither. A lone half reads as configured and does nothing (resolution requires both), so it is refused where the person setting it can see why rather than silently ignored at search time. - The key must be a Cohere-purpose secret the organization holds, checked at write time - the mirror of the embedding-key check, and for the same reason: resolution degrades a bad key to no reranking, so this is the one moment a wrong choice is visible. Unlike the embedding model, reranking can be changed after creation. Update sends the pair only when the caller actually included it (read from model_fields_set), so an update about something else leaves reranking alone and sending both as null is how it is turned off - a distinction the repo's None-means-skip convention cannot make, hence the explicit set_rerank flag. The Read schema exposes the model and the secret id; the id names a vault row, never its value. Refs #142 --- backend/app/repositories/knowledge_base.py | 13 +++ backend/app/schemas/knowledge_base.py | 25 +++++ backend/app/services/knowledge_base.py | 64 ++++++++++++ backend/tests/test_kb_scoping.py | 107 ++++++++++++++++++++- 4 files changed, 208 insertions(+), 1 deletion(-) diff --git a/backend/app/repositories/knowledge_base.py b/backend/app/repositories/knowledge_base.py index f68976bd1..d885d36e3 100644 --- a/backend/app/repositories/knowledge_base.py +++ b/backend/app/repositories/knowledge_base.py @@ -93,6 +93,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. @@ -115,6 +117,8 @@ async def create( embedding_model=embedding_model, embedding_dim=embedding_dim, 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) @@ -130,6 +134,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, ) -> KnowledgeBase: if name is not None: db_kb.name = name @@ -137,6 +144,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 await db.flush() await db.refresh(db_kb) return db_kb diff --git a/backend/app/schemas/knowledge_base.py b/backend/app/schemas/knowledge_base.py index 309afb08d..68ba47ef2 100644 --- a/backend/app/schemas/knowledge_base.py +++ b/backend/app/schemas/knowledge_base.py @@ -38,6 +38,22 @@ class KnowledgeBaseCreate(BaseSchema): "embeddings. Omit to use the deployment's key." ), ) + 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=( @@ -61,6 +77,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) class KnowledgeBaseRead(BaseSchema, TimestampSchema): @@ -81,6 +102,10 @@ class KnowledgeBaseRead(BaseSchema, TimestampSchema): embedding_model: str embedding_dim: int 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/services/knowledge_base.py b/backend/app/services/knowledge_base.py index 04f96c48e..e6b6f55e5 100644 --- a/backend/app/services/knowledge_base.py +++ b/backend/app/services/knowledge_base.py @@ -33,6 +33,7 @@ deployment_defaults, deployment_embedding, ) +from app.services.rerank_resolution import RERANK_KEY_PURPOSES logger = logging.getLogger(__name__) @@ -304,6 +305,9 @@ async def create( embedding_model, embedding_dim = chosen_embedding(data.embedding_model) if data.embedding_secret_id is not None: await self._check_embedding_secret(data.embedding_secret_id, organization_id=org_id) + self._check_rerank_pair(data.rerank_model, data.rerank_secret_id) + if data.rerank_secret_id is not None: + await self._check_rerank_secret(data.rerank_secret_id, organization_id=org_id) return await knowledge_base_repo.create( self.db, name=data.name, @@ -316,6 +320,8 @@ async def create( embedding_model=embedding_model, embedding_dim=embedding_dim, embedding_secret_id=data.embedding_secret_id, + rerank_model=data.rerank_model, + rerank_secret_id=data.rerank_secret_id, ) async def _check_embedding_secret( @@ -349,6 +355,50 @@ async def _check_embedding_secret( details={"purpose": row.purpose}, ) + @staticmethod + def _check_rerank_pair(model: str | None, secret_id: UUID | None) -> None: + """A reranker is a 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. Refused + here, where the person setting it can see why, rather than silently + ignored 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)}, + ) + + async def _check_rerank_secret(self, secret_id: UUID, *, organization_id: UUID | None) -> None: + """Refuse a rerank key the organization does not hold, or the wrong kind. + + The mirror of :meth:`_check_embedding_secret`, and checked at the same + moment and for the same reason: resolution degrades a bad key to no + reranking, so creation is the one place a wrong choice is visible. + """ + 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: + 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, @@ -362,12 +412,26 @@ 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( + data.rerank_secret_id, organization_id=kb.organization_id + ) return await knowledge_base_repo.update( self.db, db_kb=kb, 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, ) async def _usable_config( diff --git a/backend/tests/test_kb_scoping.py b/backend/tests/test_kb_scoping.py index a89db0305..f5f287839 100644 --- a/backend/tests/test_kb_scoping.py +++ b/backend/tests/test_kb_scoping.py @@ -11,7 +11,7 @@ from app.db.models.knowledge_base import KBScope, KnowledgeBase from app.db.models.resource_grant import GrantLevel, Visibility from app.repositories.rag_document import CollectionCounts -from app.schemas.knowledge_base import KnowledgeBaseCreate +from app.schemas.knowledge_base import KnowledgeBaseCreate, KnowledgeBaseUpdate from app.services.ingestion_config import deployment_defaults from app.services.knowledge_base import KnowledgeBaseService, _with_counts @@ -572,3 +572,108 @@ 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_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.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_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")), + ), + 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 From 3e2bedf65dabea7a7f7f8e36246cea9a379953d0 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Tue, 18 Aug 2026 11:19:46 +0200 Subject: [PATCH 08/37] test(rag): pin the rerank key never escapes and stays in its tenant The OpenAPI no-secret-escapes sweep found KnowledgeBaseRead.rerank_secret_id as a new credential-shaped field; it is an id naming a revocable vault reference, never the key, exactly like embedding_secret_id, so it joins the allowlist with that reason rather than being removed from the response. Adds a resolution test that the rerank key is looked up scoped to the collection's own organization - the tenant boundary at the resolution layer, on top of the vault's own cross-tenant refusal (test_vault.py) and the org-scoped secret query. Refs #142 --- backend/tests/api/test_no_secret_escapes.py | 4 ++++ backend/tests/test_rerank_resolution.py | 9 +++++++++ 2 files changed, 13 insertions(+) diff --git a/backend/tests/api/test_no_secret_escapes.py b/backend/tests/api/test_no_secret_escapes.py index af6cc55ab..35d3eb1ac 100644 --- a/backend/tests/api/test_no_secret_escapes.py +++ b/backend/tests/api/test_no_secret_escapes.py @@ -101,6 +101,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/test_rerank_resolution.py b/backend/tests/test_rerank_resolution.py index 02c94f0b6..589eb600b 100644 --- a/backend/tests/test_rerank_resolution.py +++ b/backend/tests/test_rerank_resolution.py @@ -107,6 +107,15 @@ async def test_a_repr_never_carries_the_key(self): 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 TestDegradationTurnsRerankingOff: """A chosen key that is gone drops reranking to off, with a line saying why. From 217eaed660d559f53d2633a2e131fec25e98302d Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Tue, 18 Aug 2026 11:21:40 +0200 Subject: [PATCH 09/37] docs(rag): document per-collection reranking MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a "Reranking — a second pass, off unless configured" subsection under the embeddings heading in file-processing.md: what a reranker is and where it wires into retrieval, the per-KB Cohere configuration mirroring embeddings, the off-by-default design and its degradation to no-rerank, the runtime fallback to distance order, and the metered-spend divergence from genai-prices plus the /rag/search metering gap it closes. Refs #142 --- docs/file-processing.md | 43 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/docs/file-processing.md b/docs/file-processing.md index 79cea5fb7..370a6aef6 100644 --- a/docs/file-processing.md +++ b/docs/file-processing.md @@ -375,6 +375,49 @@ 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. +### 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, because an agent's bound collections share one +organization and so one reranker. + +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. | +| **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) 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. + ### Vector Storage Vectors are stored in **pgvector** using the existing PostgreSQL database. No additional services needed. From bf12b00eb473b42deaad5334eb3bbab63b771b5a Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Tue, 18 Aug 2026 12:15:01 +0200 Subject: [PATCH 10/37] feat(rag): set a collection's reranker from the KB create dialog Mirrors the embedding-key picker: a "Reranking" disclosure in CreateKBDialog with an Off/key Select over the organization's cohere-purpose vault keys (plus inline add-a-key), placed after Embeddings. There is one reranker and no endpoint listing them, so the model is a frontend constant (rerank-v3.5) and choosing a key is choosing to rerank - the submit sends the key and the model together or neither, matching the backend's both-or-neither rule. Create-time only, exactly as embeddings are: the KB detail page has no edit control for these config fields, and the backend accepts the pair at create. rerank_model / rerank_secret_id added to CreateKnowledgeBaseInput. A guided-tour stop (flow-kb-field-rerank) mirrors the embeddings one, gated on collections:edit and anchored on the disclosure. Tests cover both branches: Off posts neither field, a chosen key posts rerank_secret_id and rerank_model=rerank-v3.5. Refs #142 --- frontend/messages/en.json | 9 +++ .../kb/create-kb-dialog.integration.test.tsx | 9 +-- .../components/kb/create-kb-dialog.test.tsx | 40 +++++++++++++ .../src/components/kb/create-kb-dialog.tsx | 60 +++++++++++++++++++ frontend/src/lib/onboarding/flows.test.ts | 2 + frontend/src/lib/onboarding/flows.ts | 9 +++ frontend/src/types/knowledge-base.ts | 7 +++ 7 files changed, 132 insertions(+), 4 deletions(-) diff --git a/frontend/messages/en.json b/frontend/messages/en.json index 6844d8184..35f845dd1 100644 --- a/frontend/messages/en.json +++ b/frontend/messages/en.json @@ -1974,6 +1974,11 @@ "reasoning": "Reasoning", "remove": "Remove", "removeNamed": "Remove {name}?", + "rerank": "Reranking", + "rerankHelp": "A reranker reorders search results by relevance. Off unless a key is set.", + "rerankKey": "Reranking key", + "rerankKeyName": "Cohere (reranking)", + "rerankOff": "Off", "reusableIntegrations": "Reusable integrations", "save": "Save", "saving": "Saving…", @@ -2658,6 +2663,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." 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 07bc240ae..a66a42ce4 100644 --- a/frontend/src/components/kb/create-kb-dialog.integration.test.tsx +++ b/frontend/src/components/kb/create-kb-dialog.integration.test.tsx @@ -301,9 +301,10 @@ 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 c87f67b3b..ac01c190c 100644 --- a/frontend/src/components/kb/create-kb-dialog.test.tsx +++ b/frontend/src/components/kb/create-kb-dialog.test.tsx @@ -51,6 +51,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" }); @@ -145,6 +160,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 91bceef5f..43dcba17e 100644 --- a/frontend/src/components/kb/create-kb-dialog.tsx +++ b/frontend/src/components/kb/create-kb-dialog.tsx @@ -48,6 +48,15 @@ const EMBEDDING_KEY_PURPOSE = "openrouter"; /** Sentinel for "the deployment's key" - a Select item may not be empty. */ const DEPLOYMENT_KEY = "__deployment__"; +/** The purpose a key must carry to pay for reranking - mirrors the backend. */ +const RERANK_KEY_PURPOSE = "cohere"; + +/** Sentinel for "no reranking" - a Select item may not be empty. */ +const RERANK_OFF = "__off__"; + +/** The one reranker there is. No endpoint lists them, so it is a constant. */ +const DEFAULT_RERANK_MODEL = "rerank-v3.5"; + interface EmbeddingModels { default: string; models: { model: string; dim: number }[]; @@ -68,11 +77,13 @@ export function CreateKBDialog({ open, onOpenChange, onCreated }: CreateKBDialog const [ingestion, setIngestion] = useState(DEFAULT_INGESTION_CONFIG); const [embeddingModel, setEmbeddingModel] = 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 embeddingKeys = secrets.filter((secret) => secret.purpose === EMBEDDING_KEY_PURPOSE); + const rerankKeys = secrets.filter((secret) => secret.purpose === RERANK_KEY_PURPOSE); // Which models this build can index with. A build property, not tenant data, // so it never goes stale while a dialog is open. // @@ -105,6 +116,7 @@ export function CreateKBDialog({ open, onOpenChange, onCreated }: CreateKBDialog setIngestion(DEFAULT_INGESTION_CONFIG); setEmbeddingModel(null); setEmbeddingSecretId(null); + setRerankSecretId(null); setErrors({}); }; @@ -125,6 +137,12 @@ export function CreateKBDialog({ open, onOpenChange, onCreated }: CreateKBDialog input.embedding_model = embeddingModel; } 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); @@ -333,6 +351,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/lib/onboarding/flows.test.ts b/frontend/src/lib/onboarding/flows.test.ts index 0a7e917d1..486f093e8 100644 --- a/frontend/src/lib/onboarding/flows.test.ts +++ b/frontend/src/lib/onboarding/flows.test.ts @@ -102,6 +102,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", @@ -384,6 +385,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 3c9c35c6f..2962447b3 100644 --- a/frontend/src/lib/onboarding/flows.ts +++ b/frontend/src/lib/onboarding/flows.ts @@ -332,6 +332,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/types/knowledge-base.ts b/frontend/src/types/knowledge-base.ts index 4a50158ad..3be0d0bc0 100644 --- a/frontend/src/types/knowledge-base.ts +++ b/frontend/src/types/knowledge-base.ts @@ -147,6 +147,13 @@ export interface CreateKnowledgeBaseInput { embedding_model?: 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; } /** A single document tracked in a KB's underlying vector collection. */ From b2a84a0ab8d59206718afb9ec8256d8b92aeb963 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Tue, 18 Aug 2026 12:47:59 +0200 Subject: [PATCH 11/37] fix(rag): reuse the existing cohere purpose instead of duplicating it Standing the stack up surfaced it: `cohere` is already a model-provider purpose (pydantic-ai supports Cohere chat models, and the provider list mints a purpose for every provider), so adding a `cohere` service entry produced two purposes with the same id - all_purposes() returned it as both model_provider and other, which the vault picker would render twice. A Cohere API key reranks and chats alike, so the right fix is to reuse the existing purpose rather than mint a second: drop the services.json entry and let RERANK_KEY_PURPOSES point at the model-provider one. This corrects the earlier choice of category "other", which was made before it was known that the id was already taken. No functional change to the rerank flow: a key stored under the cohere purpose still matches RERANK_KEY_PURPOSES, and the builder's rerank picker still filters on purpose == "cohere". Refs #142 --- backend/app/core/catalog/services.json | 8 -------- backend/app/services/rerank_resolution.py | 7 ++++--- 2 files changed, 4 insertions(+), 11 deletions(-) diff --git a/backend/app/core/catalog/services.json b/backend/app/core/catalog/services.json index 781845520..b71604f25 100644 --- a/backend/app/core/catalog/services.json +++ b/backend/app/core/catalog/services.json @@ -53,13 +53,5 @@ "kind": "api_key", "help_url": "https://app.daytona.io/dashboard/keys", "description": "Cloud sandboxes an agent works in, billed to this organization's own Daytona account." - }, - { - "id": "cohere", - "label": "Cohere", - "category": "other", - "kind": "api_key", - "help_url": "https://dashboard.cohere.com/api-keys", - "description": "Rerank knowledge-search results with Cohere, billed to this organization's own key." } ] diff --git a/backend/app/services/rerank_resolution.py b/backend/app/services/rerank_resolution.py index 74d72345d..616c3883b 100644 --- a/backend/app/services/rerank_resolution.py +++ b/backend/app/services/rerank_resolution.py @@ -33,9 +33,10 @@ logger = logging.getLogger(__name__) -# Secret purposes that can pay for reranking. Cohere is the only reranker today; -# the tuple exists so a second provider is one entry, not a hunt, and so the -# `cohere` entry in the services catalog has a consumer that names it. +# 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",) From 72e0744daf65ecbab6ba33d7af093b0da408612d Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Tue, 18 Aug 2026 14:50:06 +0200 Subject: [PATCH 12/37] fix(rag): rerank on the agent-run path, not only /rag/search The knowledge capability built its own RetrievalService with no reranker resolver, so an agent searching its bound collections never reranked - only the /rag/search route did. That contradicts the contract: rerank spend is meant to be recorded on both the agent-run path and the route (#142 done-when), and the whole "a rerank during a knowledge search books automatically" framing assumed the run's open ledger would see it. Both paths now build their RetrievalService with one shared composition point, build_reranker() in reranker.py, which resolves a collection's credential and binds it to the CohereReranker. deps.get_retrieval_service and the knowledge tool's get_retrieval_service both pass it, so reranking is wired identically and a second provider is a branch in one place. The agent-run path's rerank cost books to the run's ledger, which was already open. Found auditing the branch against the issue. Tests assert both paths wire build_reranker and that it degrades to None when unconfigured. Refs #142 --- .../agents/capabilities/knowledge/_search.py | 9 +++- backend/app/api/deps.py | 23 +++------ backend/app/services/rag/reranker.py | 17 +++++++ backend/tests/test_reranker.py | 50 ++++++++++++++++++- 4 files changed, 79 insertions(+), 20 deletions(-) diff --git a/backend/app/agents/capabilities/knowledge/_search.py b/backend/app/agents/capabilities/knowledge/_search.py index cf72c84db..8bbd55070 100644 --- a/backend/app/agents/capabilities/knowledge/_search.py +++ b/backend/app/agents/capabilities/knowledge/_search.py @@ -8,6 +8,7 @@ from app.core.exceptions import AppException, ExternalServiceError from app.services.embedding_resolution import embeddings_for_collection 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 PgVectorStore @@ -30,7 +31,13 @@ def get_retrieval_service() -> "BaseRetrievalService": vector_store = PgVectorStore( rag_settings, embedding_service, resolver=embeddings_for_collection ) - _retrieval_service = RetrievalService(vector_store, 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( + vector_store, rag_settings, reranker_resolver=build_reranker + ) return _retrieval_service diff --git a/backend/app/api/deps.py b/backend/app/api/deps.py index 3ce995f94..bc4d1c5d8 100644 --- a/backend/app/api/deps.py +++ b/backend/app/api/deps.py @@ -858,11 +858,10 @@ async def verify_api_key( 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 BaseReranker, CohereReranker +from app.services.rag.reranker import build_reranker from app.services.rag.retrieval import RetrievalService from app.services.rag.vectorstore import PgVectorStore from app.services.rag.vectorstore import BaseVectorStore -from app.services.rerank_resolution import reranker_for_collection def get_embedding_service(request: Request) -> EmbeddingService: @@ -887,26 +886,16 @@ def get_vectorstore(request: Request, embedder: EmbeddingSvc) -> BaseVectorStore VectorStoreSvc = Annotated[BaseVectorStore, Depends(get_vectorstore)] -async def _reranker_for_collection(collection_name: str) -> BaseReranker | None: - """Bind a collection's resolved reranker credential to a concrete reranker. +def get_retrieval_service(vector_store: VectorStoreSvc) -> RetrievalService: + """Create RetrievalService instance. - The composition root for reranking: resolution answers whether a collection - is configured and with whose key, and this turns that into the one - implementation there is. A second provider is a branch here, not a change to - retrieval. + The reranker resolver is `build_reranker`, the one composition point shared + with the agent-run knowledge tool, so both paths rerank the same way. """ - resolved = await reranker_for_collection(collection_name) - if resolved is None: - return None - return CohereReranker(model=resolved.model, api_key=resolved.api_key) - - -def get_retrieval_service(vector_store: VectorStoreSvc) -> RetrievalService: - """Create RetrievalService instance.""" return RetrievalService( vector_store=vector_store, settings=settings.rag, - reranker_resolver=_reranker_for_collection, + reranker_resolver=build_reranker, ) diff --git a/backend/app/services/rag/reranker.py b/backend/app/services/rag/reranker.py index ebf4c16ab..6e70d686a 100644 --- a/backend/app/services/rag/reranker.py +++ b/backend/app/services/rag/reranker.py @@ -31,6 +31,7 @@ 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 @@ -120,3 +121,19 @@ def _spend_entry(self, document_count: int) -> SpendEntry: cost_usd=_PRICE_PER_SEARCH_UNIT_USD * units, priced=True, ) + + +async def build_reranker(collection_name: str) -> 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. + """ + resolved = await reranker_for_collection(collection_name) + if resolved is None: + return None + return CohereReranker(model=resolved.model, api_key=resolved.api_key) diff --git a/backend/tests/test_reranker.py b/backend/tests/test_reranker.py index 46e15146f..7be7e90fe 100644 --- a/backend/tests/test_reranker.py +++ b/backend/tests/test_reranker.py @@ -10,13 +10,14 @@ from decimal import Decimal from types import SimpleNamespace -from unittest.mock import AsyncMock +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 +from app.services.rag.reranker import CohereReranker, build_reranker +from app.services.rerank_resolution import ResolvedReranker pytestmark = pytest.mark.anyio @@ -108,3 +109,48 @@ def test_the_client_is_built_lazily_from_the_key(self): reranker = CohereReranker(model="rerank-v3.5", api_key="co-key") assert reranker._client is None assert reranker.client is not None + + +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") 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") + 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 From a1ced174063b28e206ff95dff5b8b37962072abb Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Tue, 18 Aug 2026 15:05:04 +0200 Subject: [PATCH 13/37] docs(rag): state the first-collection rerank rule and the min_score scale Two honesty fixes surfaced in review, both docstring/doc only - no behaviour change. retrieve_multi resolves the reranker from collection_names[0] and runs it once over the union. The docstring justified that with "an agent's bound collections share one organization and so one reranker", which is true on the agent-run path but silent about /rag/search, where a caller may pass any readable set of one organization and the first collection's setting governs the whole union - a set led by a plain collection stays in distance order even if a later one reranks. Said plainly now, in the docstring and in docs/file-processing.md. min_score gates recall on the vector-distance score; after reranking, score carries the reranker's relevance judgement on a different scale and min_score is not re-applied. A caller thresholding on a reranked result's score as if it were the recall score would be wrong. Noted on _recall. --- backend/app/services/rag/retrieval.py | 21 ++++++++++++++++----- docs/file-processing.md | 7 +++++-- 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/backend/app/services/rag/retrieval.py b/backend/app/services/rag/retrieval.py index 83726fb2d..7791ca626 100644 --- a/backend/app/services/rag/retrieval.py +++ b/backend/app/services/rag/retrieval.py @@ -189,6 +189,11 @@ async def _recall( 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'", @@ -290,11 +295,17 @@ async def retrieve_multi( 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: the bound collections of - one agent share one organization and so one reranker, and reranking each - collection separately then merging the winners would rank against the - wrong pool. Absent a reranker this is byte-for-byte the previous merge - - each collection's top `limit`, fused, sorted, deduplicated, truncated. + collection's candidates, not per collection: reranking each collection + separately then merging the winners would rank against the wrong pool. + The *first* collection's configuration decides which reranker, if any. + On the agent-run path that is unambiguous - an agent's bound collections + share one organization and one configuration - but `/rag/search` may pass + any readable set of one organization, and there the first collection + governs the whole union: a set led by a reranking collection pays for and + reorders all of them, one led by a plain collection leaves them in + distance order. Absent a reranker this is byte-for-byte the previous + merge - each collection's top `limit`, fused, sorted, deduplicated, + truncated. """ reranker = await self._reranker_for(collection_names[0]) if collection_names else None multiplier = _RERANK_FETCH_MULTIPLIER if reranker else _DEFAULT_FETCH_MULTIPLIER diff --git a/docs/file-processing.md b/docs/file-processing.md index 370a6aef6..4b5371892 100644 --- a/docs/file-processing.md +++ b/docs/file-processing.md @@ -383,8 +383,11 @@ 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, because an agent's bound collections share one -organization and so one reranker. +every collection's candidates. Which reranker is the *first* collection's: 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, and a union led by a plain collection is left in distance +order even if a later one reranks. Configured **per knowledge base**, mirroring embeddings, by `app/services/rerank_resolution.py`: From d97c2bb08a27620c73ce60c9c3ea27a872a44706 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Tue, 18 Aug 2026 15:48:17 +0200 Subject: [PATCH 14/37] feat(rag): edit or turn off a collection's reranker after creation The backend already let a knowledge base's reranker be changed - the update schema reads rerank_model and rerank_secret_id as a pair, sets them when both are sent and turns reranking off when both are null - and the Read schema and docs both promise "unlike the embedding model this can be changed later". Only the UI disagreed: reranking could be set at creation and never touched again, so the whole change/turn-off path was reachable only through the API. An unreachable feature is worse than none (#61), so the console now matches what the backend and the docs already say. The detail page gains a Reranking panel beside the ingestion one - its own section, because reranking is a retrieval-time setting that changes on a different day than how the documents were read - stating whether searches are reranked and with which key, resolved from the vault and falling back to a neutral label for a reader who cannot list secrets. Its Edit opens a dialog that mirrors the ingestion one: a key picker whose "off" is one of its options, an inline "add a Cohere key", and a Save that sends the pair the backend reads together (a key with the one model, or two nulls). `updateRerank` on useKBDetail is the mutation, patching the same /kb/{id}. Also: a kb-rerank tour stop so the new section is not invisible to the walkthrough; the three rerank constants moved to src/lib/rerank-config.ts so the create and edit dialogs cannot drift; rerank_model and rerank_secret_id added to the KnowledgeBase read type (and the nine fixtures that build one). Verified: new components and the hook branch fully covered; make lint-frontend and the coverage gate green (the one miss is the pre-existing resume.ts:45 artifact, covered on CI). --- frontend/messages/en.json | 11 + .../rag/[id]/counts.integration.test.tsx | 2 + .../delete-collection.integration.test.tsx | 2 + .../kb-detail-sections.integration.test.tsx | 2 + .../[locale]/(dashboard)/rag/[id]/page.tsx | 18 ++ .../rag/[id]/status.integration.test.tsx | 2 + .../(dashboard)/rag/copy.integration.test.tsx | 2 + .../rag/rag-page.integration.test.tsx | 2 + .../agents/collection-picker.test.tsx | 2 + .../src/components/kb/create-kb-dialog.tsx | 10 +- frontend/src/components/kb/index.ts | 2 + .../src/components/kb/rerank-dialog.test.tsx | 227 ++++++++++++++++++ frontend/src/components/kb/rerank-dialog.tsx | 155 ++++++++++++ .../src/components/kb/rerank-panel.test.tsx | 104 ++++++++ frontend/src/components/kb/rerank-panel.tsx | 72 ++++++ ...reusable-integrations.integration.test.tsx | 2 + .../src/hooks/use-knowledge-bases.test.tsx | 42 ++++ frontend/src/hooks/use-knowledge-bases.ts | 26 ++ .../hooks/use-reusable-integrations.test.tsx | 2 + frontend/src/lib/onboarding/tour.test.ts | 2 +- frontend/src/lib/onboarding/tour.ts | 1 + frontend/src/lib/rerank-config.ts | 20 ++ frontend/src/types/knowledge-base.ts | 22 ++ 23 files changed, 720 insertions(+), 10 deletions(-) create mode 100644 frontend/src/components/kb/rerank-dialog.test.tsx create mode 100644 frontend/src/components/kb/rerank-dialog.tsx create mode 100644 frontend/src/components/kb/rerank-panel.test.tsx create mode 100644 frontend/src/components/kb/rerank-panel.tsx create mode 100644 frontend/src/lib/rerank-config.ts diff --git a/frontend/messages/en.json b/frontend/messages/en.json index 35f845dd1..4b8a25861 100644 --- a/frontend/messages/en.json +++ b/frontend/messages/en.json @@ -1975,10 +1975,16 @@ "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…", @@ -2023,6 +2029,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", @@ -2855,6 +2862,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/src/app/[locale]/(dashboard)/rag/[id]/counts.integration.test.tsx b/frontend/src/app/[locale]/(dashboard)/rag/[id]/counts.integration.test.tsx index 0b9725834..1a655be4a 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 @@ -51,6 +51,8 @@ const COLLECTION: KnowledgeBase = { ingestion_config: DEFAULT_INGESTION_CONFIG, embedding_model: "text-embedding-3-small", 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 a689b0cf8..cded3b35c 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 @@ -62,6 +62,8 @@ const COLLECTION: KnowledgeBase = { ingestion_config: DEFAULT_INGESTION_CONFIG, embedding_model: "text-embedding-3-small", 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 377d810f8..6dc833126 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 @@ -60,6 +60,8 @@ const KB: KnowledgeBase = { ingestion_config: DEFAULT_INGESTION_CONFIG, embedding_model: "text-embedding-3-large", 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 33f80f465..d0161177d 100644 --- a/frontend/src/app/[locale]/(dashboard)/rag/[id]/page.tsx +++ b/frontend/src/app/[locale]/(dashboard)/rag/[id]/page.tsx @@ -16,6 +16,8 @@ import { SyncSourcesSection } from "@/components/rag/sync-sources-section"; import { FileViewer } from "@/components/kb/file-viewer"; 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 } from "@/hooks"; import { overrideSize } from "@/lib/ingestion-config"; @@ -50,6 +52,7 @@ export default function KBDetailPage({ params }: KBDetailPageProps) { refresh, loadMoreDocuments, updateIngestion, + updateRerank, uploadDocument, deleteDocument, deleteCollection, @@ -69,6 +72,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 [overrideOpen, setOverrideOpen] = useState(false); /** * What a destructive control has asked for and not yet been granted. @@ -239,6 +243,12 @@ export default function KBDetailPage({ params }: KBDetailPageProps) { setIngestionOpen(true) : undefined} /> + {/* After ingestion: reranking is the other per-collection retrieval knob, + and the only one changeable after creation. */} +
+ setRerankOpen(true) : undefined} /> +
+ + + = {}): KnowledgeBase { ingestion_config: DEFAULT_INGESTION_CONFIG, embedding_model: "text-embedding-3-large", 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/kb/create-kb-dialog.tsx b/frontend/src/components/kb/create-kb-dialog.tsx index 43dcba17e..b61811514 100644 --- a/frontend/src/components/kb/create-kb-dialog.tsx +++ b/frontend/src/components/kb/create-kb-dialog.tsx @@ -35,6 +35,7 @@ import { ingestionProblems, sameIngestion, } from "@/lib/ingestion-config"; +import { DEFAULT_RERANK_MODEL, RERANK_KEY_PURPOSE, RERANK_OFF } from "@/lib/rerank-config"; import type { CreateKnowledgeBaseInput, IngestionConfig, KBScope } from "@/types"; import { useTranslations } from "next-intl"; @@ -48,15 +49,6 @@ const EMBEDDING_KEY_PURPOSE = "openrouter"; /** Sentinel for "the deployment's key" - a Select item may not be empty. */ const DEPLOYMENT_KEY = "__deployment__"; -/** The purpose a key must carry to pay for reranking - mirrors the backend. */ -const RERANK_KEY_PURPOSE = "cohere"; - -/** Sentinel for "no reranking" - a Select item may not be empty. */ -const RERANK_OFF = "__off__"; - -/** The one reranker there is. No endpoint lists them, so it is a constant. */ -const DEFAULT_RERANK_MODEL = "rerank-v3.5"; - interface EmbeddingModels { default: string; models: { model: string; dim: number }[]; 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..1b26f4a3a --- /dev/null +++ b/frontend/src/components/kb/rerank-panel.test.tsx @@ -0,0 +1,104 @@ +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, + 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 090d7f08d..4f5121d29 100644 --- a/frontend/src/components/kb/reusable-integrations.integration.test.tsx +++ b/frontend/src/components/kb/reusable-integrations.integration.test.tsx @@ -55,6 +55,8 @@ function kb(id: string, name: string, collection: string): KnowledgeBase { ingestion_config: DEFAULT_INGESTION_CONFIG, embedding_model: "text-embedding-3-large", 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 c409fe10c..8972e301d 100644 --- a/frontend/src/hooks/use-knowledge-bases.test.tsx +++ b/frontend/src/hooks/use-knowledge-bases.test.tsx @@ -538,6 +538,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 0000bd0fb..47bcbca01 100644 --- a/frontend/src/hooks/use-knowledge-bases.ts +++ b/frontend/src/hooks/use-knowledge-bases.ts @@ -28,6 +28,7 @@ import type { KBDocumentList, KnowledgeBase, KnowledgeBaseList, + UpdateRerankInput, } from "@/types"; export function useKnowledgeBases() { @@ -333,6 +334,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], + ); + const uploadDocument = useCallback( async (file: File, override?: IngestionOverride) => { if (!id) return; @@ -588,6 +613,7 @@ export function useKBDetail(id: string | null) { refresh, loadMoreDocuments, updateIngestion, + updateRerank, uploadDocument, deleteDocument, deleteCollection, diff --git a/frontend/src/hooks/use-reusable-integrations.test.tsx b/frontend/src/hooks/use-reusable-integrations.test.tsx index 3a8611ed4..8bd62f91c 100644 --- a/frontend/src/hooks/use-reusable-integrations.test.tsx +++ b/frontend/src/hooks/use-reusable-integrations.test.tsx @@ -54,6 +54,8 @@ const TARGET: KnowledgeBase = { ingestion_config: DEFAULT_INGESTION_CONFIG, embedding_model: "text-embedding-3-large", 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/tour.test.ts b/frontend/src/lib/onboarding/tour.test.ts index dd1a099e1..c17cac28b 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 ad2d81540..6c9e8ff68 100644 --- a/frontend/src/lib/onboarding/tour.ts +++ b/frontend/src/lib/onboarding/tour.ts @@ -307,6 +307,7 @@ export const TOUR_STEPS: readonly TourStep[] = [ inTour: true, }, { id: "kb-ingestion", page: KB_DETAIL, target: "kb-ingestion", permission: Perm.collectionsView }, + { id: "kb-rerank", page: KB_DETAIL, target: "kb-rerank", permission: Perm.collectionsView }, { id: "kb-sync", page: KB_DETAIL, target: "kb-sync", permission: Perm.collectionsView }, { id: "orgs-new", page: ROUTES.ORGS, target: "orgs-new" }, 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 3be0d0bc0..5fae45c83 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; created_at: string; updated_at: string | null; /** @@ -156,6 +164,20 @@ export interface CreateKnowledgeBaseInput { 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; +} + /** A single document tracked in a KB's underlying vector collection. */ export interface KBDocument { id: string; From 117e8f7d31be7df6c9a7c6f90f339835cd62994f Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Tue, 18 Aug 2026 20:47:05 +0200 Subject: [PATCH 15/37] fix(rag): check the caller may use a rerank key before binding it Binding a rerank key is lending it: reranking spends it for everyone who can search the collection. `_check_rerank_secret` only looked the key up scoped to the organization, so a `collections:edit` holder who supplied the UUID of another member's private Cohere secret bound a key that `secrets:view` would refuse them - the picker never offers it, but the API takes an id and an id is guessable. Now it runs `resolve_access(..., SECRETS_VIEW)` on the row exactly as agent secret bindings do, and refuses a key the caller cannot reach as "not in this organization's vault" - the same answer as a genuine miss, so a refusal cannot be told apart and used to enumerate the vault. The mirror `_check_embedding_secret` has the identical gap; it predates this branch and is filed separately rather than folded in. --- backend/app/services/knowledge_base.py | 34 +++++++++++++++++-------- backend/tests/test_kb_scoping.py | 35 ++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 10 deletions(-) diff --git a/backend/app/services/knowledge_base.py b/backend/app/services/knowledge_base.py index e6b6f55e5..93263aa94 100644 --- a/backend/app/services/knowledge_base.py +++ b/backend/app/services/knowledge_base.py @@ -23,7 +23,7 @@ KnowledgeBaseRead, KnowledgeBaseUpdate, ) -from app.services.access import COLLECTION, visible_resource_ids +from app.services.access import COLLECTION, SECRET, resolve_access, visible_resource_ids from app.services.collection_access import CollectionAccessService, readable_kb, writable_kb from app.services.embedding_resolution import EMBEDDING_KEY_PURPOSES from app.services.ingestion_config import ( @@ -307,7 +307,7 @@ async def create( await self._check_embedding_secret(data.embedding_secret_id, organization_id=org_id) self._check_rerank_pair(data.rerank_model, data.rerank_secret_id) if data.rerank_secret_id is not None: - await self._check_rerank_secret(data.rerank_secret_id, organization_id=org_id) + 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, @@ -370,12 +370,24 @@ def _check_rerank_pair(model: str | None, secret_id: UUID | None) -> None: details={"rerank_model": model, "rerank_secret_id": str(secret_id)}, ) - async def _check_rerank_secret(self, secret_id: UUID, *, organization_id: UUID | None) -> None: - """Refuse a rerank key the organization does not hold, or the wrong kind. - - The mirror of :meth:`_check_embedding_secret`, and checked at the same - moment and for the same reason: resolution degrades a bad key to no - reranking, so creation is the one place a wrong choice is visible. + 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( @@ -385,7 +397,9 @@ async def _check_rerank_secret(self, secret_id: UUID, *, organization_id: UUID | row = await organization_secret_repo.get( self.db, secret_id, organization_id=organization_id ) - if row is None: + 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)}, @@ -421,7 +435,7 @@ async def update( self._check_rerank_pair(data.rerank_model, data.rerank_secret_id) if data.rerank_secret_id is not None: await self._check_rerank_secret( - data.rerank_secret_id, organization_id=kb.organization_id + ctx, data.rerank_secret_id, organization_id=kb.organization_id ) return await knowledge_base_repo.update( self.db, diff --git a/backend/tests/test_kb_scoping.py b/backend/tests/test_kb_scoping.py index f5f287839..e04562bf3 100644 --- a/backend/tests/test_kb_scoping.py +++ b/backend/tests/test_kb_scoping.py @@ -619,6 +619,10 @@ async def test_a_configured_pair_is_written_through(self, mock_db, unclaimed_col "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()), @@ -629,6 +633,33 @@ async def test_a_configured_pair_is_written_through(self, mock_db, unclaimed_col 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( @@ -643,6 +674,10 @@ async def test_a_key_of_the_wrong_purpose_is_refused(self, mock_db, unclaimed_co "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()) From b3bf3b1f9962f33a0f8c8709ebb57731b4264c9b Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Tue, 18 Aug 2026 20:47:28 +0200 Subject: [PATCH 16/37] fix(rag): warn when a deleted key leaves a half-configured reranker The rerank pair is written together, but deleting the chosen Cohere secret nulls rerank_secret_id through the foreign key while leaving rerank_model set. `_resolve_reranker` classified that half state as the normal, silent NOT_CONFIGURED case, so reranking stopped with no signal at all - an operator had no way to see a key deletion had quietly turned it off. Now only the genuine null/null state is silent; a model with no key resolves to SECRET_MISSING and logs the same warning the other degraded cases do. --- backend/app/services/rerank_resolution.py | 9 ++++++++- backend/tests/test_rerank_resolution.py | 10 ++++++++-- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/backend/app/services/rerank_resolution.py b/backend/app/services/rerank_resolution.py index 616c3883b..870407f1e 100644 --- a/backend/app/services/rerank_resolution.py +++ b/backend/app/services/rerank_resolution.py @@ -124,8 +124,15 @@ async def _resolve_reranker( model = kb.rerank_model secret_id = kb.rerank_secret_id organization_id = kb.organization_id - if model is None or secret_id is None or organization_id is None: + 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: diff --git a/backend/tests/test_rerank_resolution.py b/backend/tests/test_rerank_resolution.py index 589eb600b..36e33a3fa 100644 --- a/backend/tests/test_rerank_resolution.py +++ b/backend/tests/test_rerank_resolution.py @@ -87,10 +87,16 @@ async def test_a_collection_that_named_no_reranker_has_none(self): assert resolved is None secrets.get.assert_not_called() - async def test_a_model_with_no_key_is_off(self): - resolved, secrets = await _resolve(_kb(model="rerank-v3.5", secret_id=None)) + 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")) From 81b6072ba4fd1a25258da8328b149c1c302f4347 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Tue, 18 Aug 2026 20:47:32 +0200 Subject: [PATCH 17/37] fix(rag): book a failed search's spend before its transaction rolls back The query embedding is booked before the vector query it pays for, so a search that fails mid-flight has already spent - and recording that on the request session was pointless, because the failed request rolls the session back and takes the spend row with it. Failed searches therefore underreported provider cost and monthly budget usage. The failure path now books through a session of its own that commits, so the cost lands whether or not the answer did - the platform records spend even when the run fails. The success path is unchanged. --- backend/app/services/knowledge_search.py | 63 ++++++++++++++++-------- backend/tests/test_knowledge_search.py | 41 +++++++++++++++ 2 files changed, 84 insertions(+), 20 deletions(-) diff --git a/backend/app/services/knowledge_search.py b/backend/app/services/knowledge_search.py index b648d4717..1f7c68cf8 100644 --- a/backend/app/services/knowledge_search.py +++ b/backend/app/services/knowledge_search.py @@ -23,6 +23,7 @@ from typing import TYPE_CHECKING from app.agents.capabilities.budget import SpendLedger, metered_by +from app.db.session import get_db_context from app.repositories import ingestion_spend_repo if TYPE_CHECKING: @@ -59,27 +60,49 @@ async def search(self, ctx: AuthContext, request: RAGSearchRequest) -> list[Sear collections = [kb.collection_name for kb in await self.access.readable_all(ctx, names)] ledger = SpendLedger(organization_id=ctx.organization_id) - 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, - ) - 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 "", - ) - - await self._record_spend(ledger) + 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, + ) + 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 "", + ) + 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 _record_spend(self, ledger: SpendLedger) -> None: + 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 @@ -92,7 +115,7 @@ async def _record_spend(self, ledger: SpendLedger) -> None: 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( - self.db, + db, organization_id=ledger.organization_id, rag_document_id=None, model=model, diff --git a/backend/tests/test_knowledge_search.py b/backend/tests/test_knowledge_search.py index a276ad31c..d87541597 100644 --- a/backend/tests/test_knowledge_search.py +++ b/backend/tests/test_knowledge_search.py @@ -125,3 +125,44 @@ async def _retrieve(*args, **kwargs) -> list[SearchResult]: 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 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() From 0249a3f1ce639b3e161c3f279c047af7f5404fd0 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Tue, 18 Aug 2026 20:47:48 +0200 Subject: [PATCH 18/37] fix(rag): don't offer a reranking edit on an app-scoped collection An app-scoped collection carries no organization_id, so it can hold no vault key and the backend refuses one ("Only an organization collection can carry a vault key"). The detail page offered its Edit anyway, so the control could only ever fail. It now shows the Reranking panel as a read-only fact and drops Edit for app scope. The create dialog needs no change: its scope picker offers only personal and org, never app. --- frontend/src/app/[locale]/(dashboard)/rag/[id]/page.tsx | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/frontend/src/app/[locale]/(dashboard)/rag/[id]/page.tsx b/frontend/src/app/[locale]/(dashboard)/rag/[id]/page.tsx index d0161177d..f0386f072 100644 --- a/frontend/src/app/[locale]/(dashboard)/rag/[id]/page.tsx +++ b/frontend/src/app/[locale]/(dashboard)/rag/[id]/page.tsx @@ -244,9 +244,14 @@ export default function KBDetailPage({ params }: KBDetailPageProps) { {/* After ingestion: reranking is the other per-collection retrieval knob, - and the only one changeable after creation. */} + 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} /> + setRerankOpen(true) : undefined} + />
Date: Tue, 18 Aug 2026 20:48:17 +0200 Subject: [PATCH 19/37] test(e2e): scope the ingestion Edit click past the new reranking panel The detail page now carries a Reranking section with its own Edit button, so a page-wide getByRole("button", {name: "Edit"}) matches two elements and Playwright's strict mode fails the ingestion spec. Scoped to the "How documents are read" region, where the spec means to click. --- frontend/e2e/kb-ingestion.spec.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/frontend/e2e/kb-ingestion.spec.ts b/frontend/e2e/kb-ingestion.spec.ts index cd8e8a1c6..e8025d61a 100644 --- a/frontend/e2e/kb-ingestion.spec.ts +++ b/frontend/e2e/kb-ingestion.spec.ts @@ -107,7 +107,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 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(); From 17008b4a5636861a6928b4962049ad47498dc40a Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Thu, 20 Aug 2026 09:35:08 +0200 Subject: [PATCH 20/37] fix(rag): refuse a knowledge search past the organization's budget `POST /rag/search` opened a ledger and made the paid query embedding and any Cohere rerank call, recording that spend only afterward. Unlike ingestion, it never asserted the monthly cap first, so a member with collection-view access could keep spending the organization's provider keys indefinitely after the budget was exhausted. Assert the budget once access is resolved and before the metered block, the same `assert_organization_within_budget` guard ingestion uses, so an exhausted cap refuses the search before it reaches a paid call. --- backend/app/services/knowledge_search.py | 11 ++++++ backend/tests/test_knowledge_search.py | 43 +++++++++++++++++++++++- docs/file-processing.md | 6 +++- 3 files changed, 58 insertions(+), 2 deletions(-) diff --git a/backend/app/services/knowledge_search.py b/backend/app/services/knowledge_search.py index 1f7c68cf8..f378f76ef 100644 --- a/backend/app/services/knowledge_search.py +++ b/backend/app/services/knowledge_search.py @@ -25,6 +25,7 @@ from app.agents.capabilities.budget import SpendLedger, metered_by 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 @@ -55,10 +56,20 @@ async def search(self, ctx: AuthContext, request: RAGSearchRequest) -> list[Sear 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] collections = [kb.collection_name for kb in await self.access.readable_all(ctx, names)] + 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): diff --git a/backend/tests/test_knowledge_search.py b/backend/tests/test_knowledge_search.py index d87541597..8e4072193 100644 --- a/backend/tests/test_knowledge_search.py +++ b/backend/tests/test_knowledge_search.py @@ -15,7 +15,12 @@ import pytest -from app.agents.capabilities.budget import SpendEntry, book_ambient_spend +from app.agents.capabilities.budget import ( + BudgetExceeded, + BudgetScope, + SpendEntry, + book_ambient_spend, +) from app.schemas.rag import RAGSearchRequest from app.services.knowledge_search import KnowledgeSearchService from app.services.rag.models import SearchResult @@ -25,6 +30,13 @@ _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) @@ -127,6 +139,35 @@ async def _retrieve(*args, **kwargs) -> list[SearchResult]: 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 diff --git a/docs/file-processing.md b/docs/file-processing.md index 43b84ed8a..e278d0c33 100644 --- a/docs/file-processing.md +++ b/docs/file-processing.md @@ -472,7 +472,11 @@ and dated in a comment) 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. +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. From 66907bf78c22002f71d85f5b1464687785f9bfe5 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Thu, 20 Aug 2026 09:37:45 +0200 Subject: [PATCH 21/37] fix(rag): reject an unsupported rerank model at create and update An API client could pair a valid Cohere key with a typo'd or unsupported model such as `rerank-v3.5x`. The pair passed validation, the collection was stored and shown as configured, and every subsequent search sent the invalid model to Cohere - whose error `_rank_and_truncate` swallows, so reranking silently stayed off while repeatedly making a doomed request. Validate the model against SUPPORTED_RERANK_MODELS in `_check_rerank_pair`, which both create and update already call, and before the vault is read. The tuple grows by one entry when a second model is supported. --- backend/app/services/knowledge_base.py | 20 ++++++++++---- backend/app/services/rerank_resolution.py | 8 ++++++ backend/tests/test_kb_scoping.py | 32 +++++++++++++++++++++++ docs/file-processing.md | 2 +- 4 files changed, 56 insertions(+), 6 deletions(-) diff --git a/backend/app/services/knowledge_base.py b/backend/app/services/knowledge_base.py index 773e989d7..637613878 100644 --- a/backend/app/services/knowledge_base.py +++ b/backend/app/services/knowledge_base.py @@ -33,7 +33,7 @@ deployment_defaults, deployment_embedding, ) -from app.services.rerank_resolution import RERANK_KEY_PURPOSES +from app.services.rerank_resolution import RERANK_KEY_PURPOSES, SUPPORTED_RERANK_MODELS logger = logging.getLogger(__name__) @@ -368,18 +368,28 @@ async def _check_embedding_secret( @staticmethod def _check_rerank_pair(model: str | None, secret_id: UUID | None) -> None: - """A reranker is a model *and* a key, or neither. + """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. Refused - here, where the person setting it can see why, rather than silently - ignored at search time. + 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 diff --git a/backend/app/services/rerank_resolution.py b/backend/app/services/rerank_resolution.py index 870407f1e..1647619d3 100644 --- a/backend/app/services/rerank_resolution.py +++ b/backend/app/services/rerank_resolution.py @@ -39,6 +39,14 @@ # 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. diff --git a/backend/tests/test_kb_scoping.py b/backend/tests/test_kb_scoping.py index 2273c906b..47cbb511c 100644 --- a/backend/tests/test_kb_scoping.py +++ b/backend/tests/test_kb_scoping.py @@ -688,6 +688,38 @@ async def test_a_key_without_a_model_is_refused(self, mock_db, unclaimed_collect 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() diff --git a/docs/file-processing.md b/docs/file-processing.md index e278d0c33..3b55c8da3 100644 --- a/docs/file-processing.md +++ b/docs/file-processing.md @@ -447,7 +447,7 @@ Configured **per knowledge base**, mirroring embeddings, by | | | |---|---| -| **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. | +| **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 From 77770f669c4b169fe76f1746ce5121f05ecde9e5 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Thu, 20 Aug 2026 09:41:59 +0200 Subject: [PATCH 22/37] fix(rag): rerank a multi-collection union only when its collections agree `retrieve_multi` resolved the reranker from the first collection alone and applied it to the whole union. For a `/rag/search` set spanning collections with different rerank settings that sent a collection which had reranking disabled to Cohere, and billed candidates from a differently-keyed collection to the first collection's credential. Resolve the reranker for every collection and rerank only when they all agree - same model, same key (CohereReranker now compares by model and key). A mixed set - one reranking, one not, or two on different keys - falls back to the by-distance union rather than reordering on a credential that is not the collection's own. The agent-run path is unchanged: an agent's bound collections share one organization and one configuration, so they always agree. Retrieval is now safe for a mixed binding regardless of bind-time checks, so `AgentRegistryService._collection_problems` is left as is rather than made to reject a binding a user may legitimately want in distance order. --- backend/app/services/rag/reranker.py | 12 +++++ backend/app/services/rag/retrieval.py | 36 ++++++++++--- backend/tests/test_reranker.py | 23 +++++++++ backend/tests/test_retrieval_reranking.py | 61 ++++++++++++++++++++++- docs/file-processing.md | 13 +++-- 5 files changed, 131 insertions(+), 14 deletions(-) diff --git a/backend/app/services/rag/reranker.py b/backend/app/services/rag/reranker.py index 6e70d686a..c6ab95424 100644 --- a/backend/app/services/rag/reranker.py +++ b/backend/app/services/rag/reranker.py @@ -78,6 +78,18 @@ def __init__(self, model: str, api_key: str, client: AsyncClientV2 | None = None self._api_key = api_key self._client = client + 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: diff --git a/backend/app/services/rag/retrieval.py b/backend/app/services/rag/retrieval.py index 7791ca626..bf32f9c15 100644 --- a/backend/app/services/rag/retrieval.py +++ b/backend/app/services/rag/retrieval.py @@ -133,6 +133,21 @@ async def _reranker_for(self, collection_name: str) -> BaseReranker | None: return None return await self._reranker_resolver(collection_name) + @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, @@ -297,17 +312,22 @@ async def retrieve_multi( 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. - The *first* collection's configuration decides which reranker, if any. - On the agent-run path that is unambiguous - an agent's bound collections - share one organization and one configuration - but `/rag/search` may pass - any readable set of one organization, and there the first collection - governs the whole union: a set led by a reranking collection pays for and - reorders all of them, one led by a plain collection leaves them in - distance order. Absent a reranker this is byte-for-byte the previous + + 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. """ - reranker = await self._reranker_for(collection_names[0]) if collection_names else None + rerankers = [await self._reranker_for(name) for name in collection_names] + reranker = self._shared_reranker(rerankers) multiplier = _RERANK_FETCH_MULTIPLIER if reranker else _DEFAULT_FETCH_MULTIPLIER all_results: list[SearchResult] = [] diff --git a/backend/tests/test_reranker.py b/backend/tests/test_reranker.py index 7be7e90fe..fef53f3d7 100644 --- a/backend/tests/test_reranker.py +++ b/backend/tests/test_reranker.py @@ -111,6 +111,29 @@ def test_the_client_is_built_lazily_from_the_key(self): assert reranker.client is not 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.""" diff --git a/backend/tests/test_retrieval_reranking.py b/backend/tests/test_retrieval_reranking.py index 08c01dafc..c59124262 100644 --- a/backend/tests/test_retrieval_reranking.py +++ b/backend/tests/test_retrieval_reranking.py @@ -14,7 +14,7 @@ import pytest from app.services.rag.models import SearchResult -from app.services.rag.reranker import BaseReranker +from app.services.rag.reranker import BaseReranker, CohereReranker from app.services.rag.retrieval import RetrievalService pytestmark = pytest.mark.anyio @@ -57,6 +57,17 @@ def _service(store: MagicMock, reranker: BaseReranker | None) -> RetrievalServic 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) -> 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)] @@ -96,3 +107,51 @@ async def test_the_collection_stamp_survives_reranking(self): store = _store_returning(_hits("a")) results = await _service(store, _ReverseReranker()).retrieve("q", "handbook", limit=1) assert results[0].metadata["collection"] == "handbook" + + +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) + 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) + 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/docs/file-processing.md b/docs/file-processing.md index 3b55c8da3..d6bc84653 100644 --- a/docs/file-processing.md +++ b/docs/file-processing.md @@ -436,11 +436,14 @@ 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. Which reranker is the *first* collection's: 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, and a union led by a plain collection is left in distance -order even if a later one reranks. +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`: From 0be16b20b35a343fe09027279bca71d66d77a998 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Thu, 20 Aug 2026 10:06:51 +0200 Subject: [PATCH 23/37] test(e2e): await the now-async howItReads before clicking its Edit main moved the ingestion panel behind a tab (#939), making `howItReads` async - it clicks the tab, then returns the region. The reranking-scoped Edit click added on this branch still called `.getByRole` on the returned promise; await it first. --- frontend/e2e/kb-ingestion.spec.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/e2e/kb-ingestion.spec.ts b/frontend/e2e/kb-ingestion.spec.ts index 8c998c636..0860c1108 100644 --- a/frontend/e2e/kb-ingestion.spec.ts +++ b/frontend/e2e/kb-ingestion.spec.ts @@ -116,7 +116,7 @@ test.describe("Ingestion settings", () => { // nobody touched alongside the one that was. // 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 howItReads(page).getByRole("button", { name: "Edit" }).click(); + 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(); From 43483293fc3cbf0207a13d95b03f6d9e77b187ce Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Thu, 20 Aug 2026 10:07:08 +0200 Subject: [PATCH 24/37] fix(rag): report search spend as retrieval, not indexing A metered `POST /rag/search` records its embedding and rerank cost in `ingestion_spend`, the table the dashboard reports as money spent "on indexing". So an ordinary search inflated the indexing subtotal, while its Cohere and embedding cost never appeared as retrieval anywhere. Add a `source` column ('ingestion' | 'retrieval') to the table, tag the search path RETRIEVAL, and split the dashboard's window figure into `model_usd`, `ingestion_usd` and `retrieval_usd` - each summed on its own so a search is not reported as indexing. Both non-run sources still count toward the monthly budget: `sum_cost_since`, which feeds the cap, stays unfiltered while the dashboard's `sum_cost_window` narrows by source. server_default 'ingestion' backfills the rows written before the column, all of them indexing, without a data migration. --- .../versions/0044_ingestion_spend_source.py | 43 +++++++++ backend/app/db/models/ingestion_spend.py | 51 +++++++--- backend/app/repositories/ingestion_spend.py | 35 +++++-- backend/app/schemas/stats.py | 23 +++-- backend/app/services/knowledge_search.py | 4 + backend/app/services/stats.py | 41 +++++++- backend/tests/test_ingestion_spend_repo.py | 51 +++++++++- backend/tests/test_knowledge_search.py | 2 + backend/tests/test_stats.py | 34 ++++++- docs/governance.md | 31 +++--- frontend/messages/en.json | 4 +- frontend/messages/pl.json | 4 +- .../dashboard/widgets/spend.test.tsx | 95 +++++++++++++++++++ .../components/dashboard/widgets/spend.tsx | 29 +++--- frontend/src/types/stats.ts | 7 +- 15 files changed, 388 insertions(+), 66 deletions(-) create mode 100644 backend/alembic/versions/0044_ingestion_spend_source.py create mode 100644 frontend/src/components/dashboard/widgets/spend.test.tsx diff --git a/backend/alembic/versions/0044_ingestion_spend_source.py b/backend/alembic/versions/0044_ingestion_spend_source.py new file mode 100644 index 000000000..75a01294a --- /dev/null +++ b/backend/alembic/versions/0044_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: 0044_ingestion_spend_source +Revises: 0043_knowledge_base_rerank +Create Date: 2026-08-20 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa + +from alembic import op + +revision: str = "0044_ingestion_spend_source" +down_revision: str | None = "0043_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/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/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/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/knowledge_search.py b/backend/app/services/knowledge_search.py index f378f76ef..c13d71aba 100644 --- a/backend/app/services/knowledge_search.py +++ b/backend/app/services/knowledge_search.py @@ -23,6 +23,7 @@ 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 @@ -120,6 +121,8 @@ async def _record_spend(db: AsyncSession, ledger: SpendLedger) -> None: 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 @@ -134,4 +137,5 @@ async def _record_spend(db: AsyncSession, ledger: SpendLedger) -> None: 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/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/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_knowledge_search.py b/backend/tests/test_knowledge_search.py index 8e4072193..5b6f198ca 100644 --- a/backend/tests/test_knowledge_search.py +++ b/backend/tests/test_knowledge_search.py @@ -21,6 +21,7 @@ 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 @@ -109,6 +110,7 @@ async def test_spend_booked_during_the_search_is_persisted_to_the_organization(s 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")]) 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/docs/governance.md b/docs/governance.md index 75cccd159..e12e478a3 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -319,18 +319,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/frontend/messages/en.json b/frontend/messages/en.json index ea24dc0ce..4e3b5269e 100644 --- a/frontend/messages/en.json +++ b/frontend/messages/en.json @@ -1572,7 +1572,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": { diff --git a/frontend/messages/pl.json b/frontend/messages/pl.json index cff0a07a1..598b26ca7 100644 --- a/frontend/messages/pl.json +++ b/frontend/messages/pl.json @@ -679,7 +679,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/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/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[]; } From 55bce968340e1381b0ab37d1c9870f88898502a4 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Thu, 20 Aug 2026 13:04:39 +0200 Subject: [PATCH 25/37] fix(rag): re-parent rerank migrations onto main's new head Auto-merging main into this branch left two Alembic heads: main added 0043_rag_document_source_path off 0042, while this branch's rerank and ingestion-spend migrations also chained off 0042. `alembic upgrade head` then aborts with "Multiple head revisions are present", which took the test job's migration suite down (four failures in test_migrations.py). Renumber the branch's two migrations onto main's head so the chain is linear again: 0042_sync_source_secret_id -> 0043_rag_document_source_path (main) -> 0044_knowledge_base_rerank (was 0043) -> 0045_ingestion_spend_source (was 0044) Verified `alembic heads` reports a single head. The upgrade/downgrade cycle runs against Postgres in CI; it skips locally with no database. --- ...ledge_base_rerank.py => 0044_knowledge_base_rerank.py} | 8 ++++---- ...ion_spend_source.py => 0045_ingestion_spend_source.py} | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) rename backend/alembic/versions/{0043_knowledge_base_rerank.py => 0044_knowledge_base_rerank.py} (91%) rename backend/alembic/versions/{0044_ingestion_spend_source.py => 0045_ingestion_spend_source.py} (85%) diff --git a/backend/alembic/versions/0043_knowledge_base_rerank.py b/backend/alembic/versions/0044_knowledge_base_rerank.py similarity index 91% rename from backend/alembic/versions/0043_knowledge_base_rerank.py rename to backend/alembic/versions/0044_knowledge_base_rerank.py index 1d1b2ae91..15357a23d 100644 --- a/backend/alembic/versions/0043_knowledge_base_rerank.py +++ b/backend/alembic/versions/0044_knowledge_base_rerank.py @@ -14,8 +14,8 @@ the embedding key there is no deployment fallback - a reranker with no key is simply off. -Revision ID: 0043_knowledge_base_rerank -Revises: 0042_sync_source_secret_id +Revision ID: 0044_knowledge_base_rerank +Revises: 0043_rag_document_source_path Create Date: 2026-08-18 """ @@ -26,8 +26,8 @@ from alembic import op -revision: str = "0043_knowledge_base_rerank" -down_revision: str | None = "0042_sync_source_secret_id" +revision: str = "0044_knowledge_base_rerank" +down_revision: str | None = "0043_rag_document_source_path" branch_labels: str | Sequence[str] | None = None depends_on: str | Sequence[str] | None = None diff --git a/backend/alembic/versions/0044_ingestion_spend_source.py b/backend/alembic/versions/0045_ingestion_spend_source.py similarity index 85% rename from backend/alembic/versions/0044_ingestion_spend_source.py rename to backend/alembic/versions/0045_ingestion_spend_source.py index 75a01294a..d345e82b8 100644 --- a/backend/alembic/versions/0044_ingestion_spend_source.py +++ b/backend/alembic/versions/0045_ingestion_spend_source.py @@ -9,8 +9,8 @@ Every row that predates the column is indexing, so `server_default` backfills them to `'ingestion'` without a data migration. -Revision ID: 0044_ingestion_spend_source -Revises: 0043_knowledge_base_rerank +Revision ID: 0045_ingestion_spend_source +Revises: 0044_knowledge_base_rerank Create Date: 2026-08-20 """ @@ -21,8 +21,8 @@ from alembic import op -revision: str = "0044_ingestion_spend_source" -down_revision: str | None = "0043_knowledge_base_rerank" +revision: str = "0045_ingestion_spend_source" +down_revision: str | None = "0044_knowledge_base_rerank" branch_labels: str | Sequence[str] | None = None depends_on: str | Sequence[str] | None = None From 0397e86c2ad5ebd6ebc1d8dcf52bf9a974aa9ab1 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Thu, 20 Aug 2026 23:07:25 +0200 Subject: [PATCH 26/37] fix(rag): close the Cohere client each rerank builds A CohereReranker builds its async Cohere HTTP client on first use and never closed it. The retrieval service that reranks is process-wide and builds a fresh reranker every search (build_reranker runs per query), so each reranked search opened a new httpx connection pool that was never returned - sockets accumulating and a fresh handshake added to every query over the life of the process. Close a client the reranker built itself in a finally around the rerank call: the reranker is request-scoped (one instance, one rerank), so the client is too. An injected client (a test's) belongs to its caller and is left open, tracked by _owns_client. The close is best-effort - a failure to return the pool is logged, never raised - so it cannot mask the rerank's own result or exception. Chose close-after-use over caching a client per (model, key): a cache would keep a decrypted vault key resident for the process lifetime, which the vault-decrypt-per-search model deliberately avoids. The extra handshake per search is marginal beside the Cohere round trip the search already makes. Verified reranker.py at 100% coverage; new tests cover the owned-client close (on success and on failure), the injected client left open, and a close failure not masking the result. Closes the codex P2 review thread on reranker.py. --- backend/app/services/rag/reranker.py | 66 +++++++++++++++++++--------- backend/tests/test_reranker.py | 49 +++++++++++++++++++++ 2 files changed, 94 insertions(+), 21 deletions(-) diff --git a/backend/app/services/rag/reranker.py b/backend/app/services/rag/reranker.py index c6ab95424..38851cb00 100644 --- a/backend/app/services/rag/reranker.py +++ b/backend/app/services/rag/reranker.py @@ -70,13 +70,19 @@ 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. + 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 @@ -102,27 +108,45 @@ async def rerank( if not results: return [] - response = await self.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(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, + 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)), ) - for item in response.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(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, document_count: int) -> SpendEntry: units = ceil(document_count / _DOCS_PER_SEARCH_UNIT) diff --git a/backend/tests/test_reranker.py b/backend/tests/test_reranker.py index fef53f3d7..fa4a7a402 100644 --- a/backend/tests/test_reranker.py +++ b/backend/tests/test_reranker.py @@ -111,6 +111,55 @@ def test_the_client_is_built_lazily_from_the_key(self): 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. From 932d50156e75ad22f2e2ecf206b5180731b4cf1a Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Fri, 21 Aug 2026 15:43:54 +0200 Subject: [PATCH 27/37] fix(rag): scope embedding and rerank resolution to the acting tenant (#1051) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What `knowledge_bases.collection_name` is indexed but **not unique** across tenants, so two organizations can name a collection the same thing. The per-collection embedding and rerank resolvers looked a knowledge base up by name alone (`get_by_collection_name(...).first()`), so on a shared name a caller could resolve another tenant's `embedding_model`/`rerank_model` and **unseal and bill its vault key** — a cross-tenant credential leak (#913). This is the store/retrieval/ingest counterpart to the route-level isolation `CollectionAccessService` already enforces; the two resolution paths are separate, and only the route one was tenant-aware. ## How (Option 1 — thread the acting tenant, no uniqueness migration) - New `knowledge_base_repo.get_for_collection(db, name, organization_id)` — two passes: the caller's own row wins; an `app`-scoped collection (owned by no organization) is the shared fallback. `organization_id=None` (a CLI ingest, no tenant) keeps the old name-only behaviour for that path alone. - `embeddings_for_collection` / `reranker_for_collection` take the acting organization and resolve through it. Identity threaded through every caller: `build_reranker`, `PgVectorStore` (`insert_document`/`search`/`get_collection_info`/`_for_collection`/ `_ensure_collection`), `RetrievalService` (`retrieve`/`retrieve_multi`/`_recall`/ `_bm25_search`), `IngestionService`, the knowledge capability + toolset, `KnowledgeSearchService`, the `/rag` routes, the ingestion DI factory and the worker ingest flow. - `organization_id` is **positional** on the two resolvers so the injected `EmbeddingResolver`/`RerankerResolver` protocols still match without rewiring, and keyword-only on the store/retrieval methods that already had positional args. ## Verification - Unit suite green (5726 passed, 0 failed). - New integration test (`test_collection_name_tenant_isolation.py`) runs **two tenants sharing one collection name against a real database**: each resolves its own embedding model and rerank key, and a third organization resolves neither. - Both gated resolver modules (`embedding_resolution`, `rerank_resolution`) stay at 100%. - `make lint-backend` exit 0, ty 0 errors. - `docs/file-processing.md` gains the tenant-scoping guarantee beside the per-collection resolution it already documents; docs-drift clean. Based off `feat/rag-reranker` (#911), which carries both resolvers — retarget to `main` on merge. Closes #913 --- .../agents/capabilities/knowledge/_search.py | 18 ++- .../agents/capabilities/knowledge/_toolset.py | 3 + backend/app/api/deps.py | 9 +- backend/app/api/routes/v1/rag.py | 6 +- backend/app/commands/rag.py | 10 +- backend/app/repositories/knowledge_base.py | 20 +++ backend/app/services/embedding_resolution.py | 15 ++- backend/app/services/knowledge_search.py | 2 + backend/app/services/rag/ingestion.py | 7 ++ backend/app/services/rag/reranker.py | 5 +- backend/app/services/rag/retrieval.py | 51 ++++++-- backend/app/services/rag/vectorstore.py | 57 ++++++--- backend/app/services/rag_document.py | 2 +- backend/app/services/rerank_resolution.py | 12 +- backend/app/worker/tasks/rag_tasks.py | 16 +-- backend/tests/api/test_error_envelope.py | 8 +- .../integration/test_chunk_insert_batching.py | 8 +- .../test_collection_name_tenant_isolation.py | 119 ++++++++++++++++++ .../tests/integration/test_platform_flows.py | 2 +- .../test_vector_store_reserved_names.py | 4 +- backend/tests/test_capability_edges.py | 30 +++-- backend/tests/test_chunk_insert_batching.py | 16 +-- backend/tests/test_collection_name_rules.py | 2 +- backend/tests/test_embedding_resolution.py | 4 +- backend/tests/test_ingestion_embedding_key.py | 22 ++-- backend/tests/test_knowledge_base_repo.py | 67 ++++++++++ backend/tests/test_rag_chunk_count.py | 1 + backend/tests/test_rag_document_lookup.py | 12 +- backend/tests/test_rag_failure_messages.py | 2 +- backend/tests/test_rerank_resolution.py | 4 +- backend/tests/test_reranker.py | 4 +- .../tests/test_reserved_collection_names.py | 2 +- backend/tests/test_retrieval_reranking.py | 30 +++-- docs/file-processing.md | 9 ++ 34 files changed, 467 insertions(+), 112 deletions(-) create mode 100644 backend/tests/integration/test_collection_name_tenant_isolation.py create mode 100644 backend/tests/test_knowledge_base_repo.py diff --git a/backend/app/agents/capabilities/knowledge/_search.py b/backend/app/agents/capabilities/knowledge/_search.py index 49f5dbbdd..8292d9f57 100644 --- a/backend/app/agents/capabilities/knowledge/_search.py +++ b/backend/app/agents/capabilities/knowledge/_search.py @@ -3,6 +3,7 @@ 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 @@ -100,6 +101,8 @@ async def search_knowledge_base( query: str, kb_collection_names: list[str] | None = None, top_k: int = 5, + *, + organization_id: UUID | None, ) -> str: """Search the knowledge base and return formatted results. @@ -109,6 +112,9 @@ 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). """ resolved = kb_collection_names if kb_collection_names else (_active_kb_collections.get() or []) if not resolved: @@ -118,10 +124,18 @@ async def search_knowledge_base( 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, + ) 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, ) 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 08172d261..4d84b4612 100644 --- a/backend/app/agents/capabilities/knowledge/_toolset.py +++ b/backend/app/agents/capabilities/knowledge/_toolset.py @@ -40,6 +40,9 @@ async def search_documents( # model chooses *what* to search, never *where*. kb_collection_names=ctx.deps.kb_collection_names, 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 as exc: raise ModelRetry("Knowledge base temporarily unavailable, please try again.") from exc diff --git a/backend/app/api/deps.py b/backend/app/api/deps.py index c259dbef6..c6d1d893c 100644 --- a/backend/app/api/deps.py +++ b/backend/app/api/deps.py @@ -940,9 +940,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 d9b230b1f..82edc25e7 100644 --- a/backend/app/api/routes/v1/rag.py +++ b/backend/app/api/routes/v1/rag.py @@ -159,7 +159,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 ) @@ -216,7 +216,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( diff --git a/backend/app/commands/rag.py b/backend/app/commands/rag.py index f7c3a303c..d7f643229 100644 --- a/backend/app/commands/rag.py +++ b/backend/app/commands/rag.py @@ -73,7 +73,11 @@ def get_rag_services() -> tuple[ ) 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 @@ -93,7 +97,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}") @@ -474,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/repositories/knowledge_base.py b/backend/app/repositories/knowledge_base.py index d885d36e3..98497bc20 100644 --- a/backend/app/repositories/knowledge_base.py +++ b/backend/app/repositories/knowledge_base.py @@ -186,3 +186,23 @@ async def list_by_collection_name(db: AsyncSession, collection_name: str) -> lis .order_by(KnowledgeBase.created_at) ) return list(result.scalars().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/services/embedding_resolution.py b/backend/app/services/embedding_resolution.py index d370f3e6a..6b4b202e6 100644 --- a/backend/app/services/embedding_resolution.py +++ b/backend/app/services/embedding_resolution.py @@ -28,6 +28,7 @@ import logging from dataclasses import dataclass from enum import StrEnum +from uuid import UUID from sqlalchemy.ext.asyncio import AsyncSession @@ -132,16 +133,24 @@ def describe(self, collection_name: str) -> str: return f"collection {collection_name!r}, which embeds on {self.key_source.explanation}" -async def embeddings_for_collection(collection_name: str) -> ResolvedEmbeddings | None: - """Resolve one collection's embedding model and credential. +async def embeddings_for_collection( + collection_name: str, organization_id: UUID | None +) -> ResolvedEmbeddings | None: + """Resolve one collection's embedding model and credential, for one organization. Returns None for a collection no knowledge base claims - the store then uses its deployment defaults, which is what such collections have always gotten. Opens its own session because the store embeds from places with no request in sight: a worker mid-ingestion, a capability mid-run. + + `organization_id` is required and scopes the resolution: `collection_name` + is not unique across tenants, so resolving by name alone could return - and + unseal and bill - another organization's key (#913). The caller passes the + organization the search or ingest is acting for; `None` only where there is + no tenant (a CLI ingest). """ 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_for_collection(db, collection_name, organization_id) if kb is None: return None api_key, key_source = await _api_key_for(db, kb) diff --git a/backend/app/services/knowledge_search.py b/backend/app/services/knowledge_search.py index c13d71aba..c8b0e429d 100644 --- a/backend/app/services/knowledge_search.py +++ b/backend/app/services/knowledge_search.py @@ -80,6 +80,7 @@ async def search(self, ctx: AuthContext, request: RAGSearchRequest) -> list[Sear collection_names=collections, limit=request.limit, min_score=request.min_score, + organization_id=ctx.organization_id, ) else: results = await self.retrieval.retrieve( @@ -88,6 +89,7 @@ async def search(self, ctx: AuthContext, request: RAGSearchRequest) -> list[Sear limit=request.limit, min_score=request.min_score, filter=request.filter or "", + organization_id=ctx.organization_id, ) except Exception: # The query embedding is booked before the vector query it pays for, diff --git a/backend/app/services/rag/ingestion.py b/backend/app/services/rag/ingestion.py index 68ccd0e41..966728c7c 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 @@ -65,10 +66,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: @@ -174,6 +180,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 index 38851cb00..cc935745c 100644 --- a/backend/app/services/rag/reranker.py +++ b/backend/app/services/rag/reranker.py @@ -26,6 +26,7 @@ from decimal import Decimal from math import ceil from typing import TYPE_CHECKING +from uuid import UUID import cohere @@ -159,7 +160,7 @@ def _spend_entry(self, document_count: int) -> SpendEntry: ) -async def build_reranker(collection_name: str) -> BaseReranker | None: +async def build_reranker(collection_name: str, organization_id: UUID | None) -> BaseReranker | None: """Bind a collection's resolved rerank credential to a concrete reranker. The one composition point for reranking: resolution answers whether a @@ -169,7 +170,7 @@ async def build_reranker(collection_name: str) -> BaseReranker | None: wired the same way in both, and a second provider is a branch here rather than a change at each call site. """ - resolved = await reranker_for_collection(collection_name) + resolved = await reranker_for_collection(collection_name, organization_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 bf32f9c15..57246a0e7 100644 --- a/backend/app/services/rag/retrieval.py +++ b/backend/app/services/rag/retrieval.py @@ -5,6 +5,7 @@ import time from abc import ABC, abstractmethod from collections.abc import Awaitable, Callable +from uuid import UUID from rank_bm25 import BM25Okapi @@ -18,7 +19,7 @@ # 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], Awaitable[BaseReranker | None]] +RerankerResolver = Callable[[str, 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 @@ -43,6 +44,8 @@ async def retrieve( limit: int = 5, min_score: float = 0.0, filter: str = "", + *, + organization_id: UUID | None, ) -> list[SearchResult]: pass @@ -94,14 +97,17 @@ 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 ) -> 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, ) if not all_results: return [] @@ -123,15 +129,20 @@ async def _bm25_search( if s > 0 ] - async def _reranker_for(self, collection_name: str) -> BaseReranker | None: + async def _reranker_for( + self, collection_name: str, organization_id: UUID | 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. + + `organization_id` scopes the resolution: a shared collection name must + resolve the caller's own rerank config, never another tenant's (#913). """ if self._reranker_resolver is None: return None - return await self._reranker_resolver(collection_name) + return await self._reranker_resolver(collection_name, organization_id) @staticmethod def _shared_reranker(rerankers: list[BaseReranker | None]) -> BaseReranker | None: @@ -179,11 +190,19 @@ async def retrieve( limit: int = 5, min_score: float = 0.0, filter: str = "", + *, + organization_id: UUID | None, ) -> list[SearchResult]: - reranker = await self._reranker_for(collection_name) + reranker = await self._reranker_for(collection_name, organization_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 + query, + collection_name, + limit, + min_score, + filter, + fetch_multiplier=multiplier, + organization_id=organization_id, ) return await self._rank_and_truncate(reranker, query, candidates, limit) @@ -196,6 +215,7 @@ async def _recall( filter: str, *, fetch_multiplier: int, + organization_id: UUID | None, ) -> list[SearchResult]: """Vector (and optionally BM25) recall, filtered and deduplicated. @@ -225,6 +245,7 @@ async def _recall( query=query, filter_expr=filter, limit=limit * fetch_multiplier, + organization_id=organization_id, ) search_time = time.time() - start_time @@ -235,7 +256,9 @@ async def _recall( ) 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 + ) if bm25_results: pipeline_results = self._rrf_fuse(pipeline_results, bm25_results) logger.info("[RETRIEVAL] Hybrid search: fused %d results", len(pipeline_results)) @@ -297,6 +320,8 @@ async def retrieve_multi( collection_names: list[str], limit: int = 5, min_score: float = 0.0, + *, + organization_id: UUID | None, ) -> list[SearchResult]: """Search several collections and merge what they return. @@ -326,14 +351,20 @@ async def retrieve_multi( merge - each collection's top `limit`, fused, sorted, deduplicated, truncated. """ - rerankers = [await self._reranker_for(name) for name in collection_names] + rerankers = [await self._reranker_for(name, organization_id) for name in collection_names] reranker = self._shared_reranker(rerankers) multiplier = _RERANK_FETCH_MULTIPLIER if reranker else _DEFAULT_FETCH_MULTIPLIER all_results: list[SearchResult] = [] for name in collection_names: recalled = await self._recall( - query, name, limit, min_score, "", fetch_multiplier=multiplier + query, + name, + limit, + min_score, + "", + fetch_multiplier=multiplier, + organization_id=organization_id, ) all_results.extend(recalled if reranker else recalled[:limit]) diff --git a/backend/app/services/rag/vectorstore.py b/backend/app/services/rag/vectorstore.py index 82f1ed78a..532033585 100644 --- a/backend/app/services/rag/vectorstore.py +++ b/backend/app/services/rag/vectorstore.py @@ -2,6 +2,7 @@ import re from abc import ABC, abstractmethod 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 @@ -42,12 +43,20 @@ async def aclose(self) -> None: """ @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, ) -> list[SearchResult]: pass @@ -60,7 +69,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 @@ -99,7 +110,7 @@ async def get_document_list(self, collection_name: str) -> RAGDocumentList: total=len(docs), ) - 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 @@ -112,7 +123,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 @@ -178,7 +189,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], 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 @@ -274,7 +285,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 + ) -> tuple[EmbeddingService, int]: """The embedder and vector width this one collection uses. Cached per (collection, model, key): an `EmbeddingService` holds an @@ -290,7 +303,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) if resolved is None: return self.embedder, self.dim cache_key = (name, resolved.model, resolved.api_key) @@ -324,10 +337,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")) @@ -363,7 +376,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: @@ -383,10 +398,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) @@ -411,7 +426,13 @@ 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, ) -> list[SearchResult]: """Nearest chunks in a collection, reporting an absent one as empty. @@ -425,7 +446,7 @@ async def search( 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) query_vector = embedder.embed_query(query) # Parse the shared `parent_doc_id == ""` filter format and apply @@ -469,7 +490,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" @@ -481,7 +504,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 ced8a6e03..ce11b76c3 100644 --- a/backend/app/services/rag_document.py +++ b/backend/app/services/rag_document.py @@ -251,7 +251,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 index 1647619d3..bd472fcba 100644 --- a/backend/app/services/rerank_resolution.py +++ b/backend/app/services/rerank_resolution.py @@ -22,6 +22,7 @@ import logging from dataclasses import dataclass from enum import StrEnum +from uuid import UUID from sqlalchemy.ext.asyncio import AsyncSession @@ -100,7 +101,9 @@ def __repr__(self) -> str: return f"ResolvedReranker(model={self.model!r}, api_key='***')" -async def reranker_for_collection(collection_name: str) -> ResolvedReranker | None: +async def reranker_for_collection( + collection_name: str, organization_id: UUID | 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 @@ -108,9 +111,14 @@ async def reranker_for_collection(collection_name: str) -> ResolvedReranker | No 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. + + `organization_id` scopes the resolution: `collection_name` is not unique + across tenants, so resolving by name alone could read another organization's + rerank config and unseal its key (#913). The caller passes the organization + the search is acting for. """ 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_for_collection(db, collection_name, organization_id) if kb is None: return None resolved, source = await _resolve_reranker(db, kb) diff --git a/backend/app/worker/tasks/rag_tasks.py b/backend/app/worker/tasks/rag_tasks.py index 570de6ba6..b906881d7 100644 --- a/backend/app/worker/tasks/rag_tasks.py +++ b/backend/app/worker/tasks/rag_tasks.py @@ -82,8 +82,10 @@ 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 + ) -> ResolvedEmbeddings | None: + resolved = await embeddings_for_collection(collection_name, organization_id) if ( resolved is not None and resolved.key_source.is_degraded @@ -135,7 +137,9 @@ async def _ingestion_service_for( embedding_service=EmbeddingService(settings=rag_settings), resolver=_announcing_resolver(), ) - return IngestionService(processor=processor, vector_store=vector_store) + return IngestionService( + processor=processor, vector_store=vector_store, organization_id=organization_id + ) async def _record_embedding_spend( @@ -191,11 +195,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( 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/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..cfeb9d0fc --- /dev/null +++ b/backend/tests/integration/test_collection_name_tenant_isolation.py @@ -0,0 +1,119 @@ +"""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 diff --git a/backend/tests/integration/test_platform_flows.py b/backend/tests/integration/test_platform_flows.py index fd245aad2..3d5c25ca4 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_vector_store_reserved_names.py b/backend/tests/integration/test_vector_store_reserved_names.py index 24e393dba..bd031bcce 100644 --- a/backend/tests/integration/test_vector_store_reserved_names.py +++ b/backend/tests/integration/test_vector_store_reserved_names.py @@ -30,7 +30,7 @@ ) -async def _no_collection_of_its_own(name: str) -> None: +async def _no_collection_of_its_own(name: str, organization_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 +90,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 9f713af07..588d89b65 100644 --- a/backend/tests/test_capability_edges.py +++ b/backend/tests/test_capability_edges.py @@ -71,7 +71,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 @@ -83,7 +83,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"] @@ -100,7 +100,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"] @@ -110,7 +110,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 @@ -121,7 +123,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 @@ -132,7 +136,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 @@ -147,7 +153,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): @@ -171,7 +179,9 @@ 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"} @@ -227,7 +237,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) @@ -248,7 +258,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_embedding_resolution.py b/backend/tests/test_embedding_resolution.py index af5fca559..2fae4f67e 100644 --- a/backend/tests/test_embedding_resolution.py +++ b/backend/tests/test_embedding_resolution.py @@ -65,9 +65,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: diff --git a/backend/tests/test_ingestion_embedding_key.py b/backend/tests/test_ingestion_embedding_key.py index 688e76f2c..dde67a409 100644 --- a/backend/tests/test_ingestion_embedding_key.py +++ b/backend/tests/test_ingestion_embedding_key.py @@ -144,12 +144,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 @@ -224,11 +224,13 @@ async def test_two_collections_on_one_key_do_not_share_each_others_name(self): model=_MODEL, dim=_DIM, api_key="", key_source=EmbeddingKeySource.DEPLOYMENT ), } - store._resolver = AsyncMock(side_effect=lambda name: resolutions[name]) + store._resolver = AsyncMock( + side_effect=lambda name, organization_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"]) @@ -319,7 +321,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): @@ -348,7 +350,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() @@ -366,13 +368,13 @@ 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: 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", @@ -393,8 +395,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_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_rag_chunk_count.py b/backend/tests/test_rag_chunk_count.py index 47d98beb8..9fd213087 100644 --- a/backend/tests/test_rag_chunk_count.py +++ b/backend/tests/test_rag_chunk_count.py @@ -57,6 +57,7 @@ def _service(processor: MagicMock) -> IngestionService: return IngestionService( processor=processor, vector_store=MagicMock(insert_document=AsyncMock(), delete_document=AsyncMock()), + organization_id=None, ) diff --git a/backend/tests/test_rag_document_lookup.py b/backend/tests/test_rag_document_lookup.py index 1715a068e..fe89f8747 100644 --- a/backend/tests/test_rag_document_lookup.py +++ b/backend/tests/test_rag_document_lookup.py @@ -42,7 +42,7 @@ def _service(docs: list[DocumentInfo]) -> IngestionService: store = MagicMock(get_documents=AsyncMock(return_value=docs)) - return IngestionService(processor=MagicMock(), vector_store=store) + return IngestionService(processor=MagicMock(), vector_store=store, organization_id=None) def _doc( @@ -131,7 +131,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 = MagicMock(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() @@ -219,7 +219,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") @@ -246,7 +246,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"), @@ -294,7 +294,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 04c6eb3d2..593e9a844 100644 --- a/backend/tests/test_rag_failure_messages.py +++ b/backend/tests/test_rag_failure_messages.py @@ -148,7 +148,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 index 36e33a3fa..07009eb10 100644 --- a/backend/tests/test_rerank_resolution.py +++ b/backend/tests/test_rerank_resolution.py @@ -70,9 +70,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 reranker_for_collection("handbook"), secrets + return await reranker_for_collection("handbook", organization_id=uuid.uuid4()), secrets class TestResolution: diff --git a/backend/tests/test_reranker.py b/backend/tests/test_reranker.py index fa4a7a402..dc417a8e3 100644 --- a/backend/tests/test_reranker.py +++ b/backend/tests/test_reranker.py @@ -191,14 +191,14 @@ async def test_an_unconfigured_collection_gets_no_reranker(self): "app.services.rag.reranker.reranker_for_collection", new=AsyncMock(return_value=None), ): - assert await build_reranker("handbook") is 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") + reranker = await build_reranker("handbook", None) assert isinstance(reranker, CohereReranker) assert reranker.model == "rerank-v3.5" 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 index c59124262..7aa0bb725 100644 --- a/backend/tests/test_retrieval_reranking.py +++ b/backend/tests/test_retrieval_reranking.py @@ -62,7 +62,7 @@ def _service_by_name(store: MagicMock, mapping: dict[str, BaseReranker | None]) settings = MagicMock() settings.enable_hybrid_search = False - async def resolver(name: str) -> BaseReranker | None: + async def resolver(name: str, organization_id: object = None) -> BaseReranker | None: return mapping.get(name) return RetrievalService(vector_store=store, settings=settings, reranker_resolver=resolver) @@ -75,22 +75,28 @@ def _hits(*names: str) -> list[SearchResult]: 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) + 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) + 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) + 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) + results = await _service(store, _FailingReranker()).retrieve( + "q", "kb", limit=2, organization_id=None + ) assert [r.content for r in results] == ["a", "b"] @@ -99,13 +105,15 @@ 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 + "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) + results = await _service(store, _ReverseReranker()).retrieve( + "q", "handbook", limit=1, organization_id=None + ) assert results[0].metadata["collection"] == "handbook" @@ -121,7 +129,9 @@ 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) + await svc.retrieve_multi( + "q", collection_names=["kb_a", "kb_b"], limit=3, organization_id=None + ) assert r1.calls == 0 assert r2.calls == 0 @@ -129,7 +139,9 @@ 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) + await svc.retrieve_multi( + "q", collection_names=["kb_a", "kb_b"], limit=3, organization_id=None + ) assert r.calls == 0 diff --git a/docs/file-processing.md b/docs/file-processing.md index 67a103572..530e0cba2 100644 --- a/docs/file-processing.md +++ b/docs/file-processing.md @@ -428,6 +428,15 @@ 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. + ### Reranking — a second pass, off unless configured Vector search orders results by embedding distance, which is a proxy for From a5f3a7c61bb1de3eba4c7602598708814627a181 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Fri, 21 Aug 2026 22:14:44 +0200 Subject: [PATCH 28/37] fix(rag): renumber the rerank migrations onto main's 0045 Main grew 0044_agent_embed_key_version and 0045_audit_impersonator after the last re-parent, so on the PR's merge ref the chain forked at 0043 with two heads: every test in tests/test_migrations.py failed with "Multiple head revisions", and e2e died in its alembic upgrade before any spec ran. Same fix as 55bce968, one head later: 0044_knowledge_base_rerank becomes 0046 (down 0045_audit_impersonator) and 0045_ingestion_spend_source becomes 0047. Verified against a real pgvector 16: alembic heads reports one head, tests/test_migrations.py passes whole (upgrade, downgrade, cycle, current-matches-head), tests/test_migration_chain.py passes, and the rerank suites (45 tests) still pass on the merged tree. Refs #911 --- ...ledge_base_rerank.py => 0046_knowledge_base_rerank.py} | 8 ++++---- ...ion_spend_source.py => 0047_ingestion_spend_source.py} | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) rename backend/alembic/versions/{0044_knowledge_base_rerank.py => 0046_knowledge_base_rerank.py} (91%) rename backend/alembic/versions/{0045_ingestion_spend_source.py => 0047_ingestion_spend_source.py} (85%) diff --git a/backend/alembic/versions/0044_knowledge_base_rerank.py b/backend/alembic/versions/0046_knowledge_base_rerank.py similarity index 91% rename from backend/alembic/versions/0044_knowledge_base_rerank.py rename to backend/alembic/versions/0046_knowledge_base_rerank.py index 15357a23d..65bde54f5 100644 --- a/backend/alembic/versions/0044_knowledge_base_rerank.py +++ b/backend/alembic/versions/0046_knowledge_base_rerank.py @@ -14,8 +14,8 @@ the embedding key there is no deployment fallback - a reranker with no key is simply off. -Revision ID: 0044_knowledge_base_rerank -Revises: 0043_rag_document_source_path +Revision ID: 0046_knowledge_base_rerank +Revises: 0045_audit_impersonator Create Date: 2026-08-18 """ @@ -26,8 +26,8 @@ from alembic import op -revision: str = "0044_knowledge_base_rerank" -down_revision: str | None = "0043_rag_document_source_path" +revision: str = "0046_knowledge_base_rerank" +down_revision: str | None = "0045_audit_impersonator" branch_labels: str | Sequence[str] | None = None depends_on: str | Sequence[str] | None = None diff --git a/backend/alembic/versions/0045_ingestion_spend_source.py b/backend/alembic/versions/0047_ingestion_spend_source.py similarity index 85% rename from backend/alembic/versions/0045_ingestion_spend_source.py rename to backend/alembic/versions/0047_ingestion_spend_source.py index d345e82b8..66f71be27 100644 --- a/backend/alembic/versions/0045_ingestion_spend_source.py +++ b/backend/alembic/versions/0047_ingestion_spend_source.py @@ -9,8 +9,8 @@ Every row that predates the column is indexing, so `server_default` backfills them to `'ingestion'` without a data migration. -Revision ID: 0045_ingestion_spend_source -Revises: 0044_knowledge_base_rerank +Revision ID: 0047_ingestion_spend_source +Revises: 0046_knowledge_base_rerank Create Date: 2026-08-20 """ @@ -21,8 +21,8 @@ from alembic import op -revision: str = "0045_ingestion_spend_source" -down_revision: str | None = "0044_knowledge_base_rerank" +revision: str = "0047_ingestion_spend_source" +down_revision: str | None = "0046_knowledge_base_rerank" branch_labels: str | Sequence[str] | None = None depends_on: str | Sequence[str] | None = None From d3b2a6da2f7f440675a2828dee7c0400770813a0 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Fri, 21 Aug 2026 22:47:04 +0200 Subject: [PATCH 29/37] fix(rag): pass organization_id from the rag-search CLI Scoping reranker resolution to the acting tenant made `retrieve`'s `organization_id` a required keyword-only argument. The `rag-search` CLI command still called `retrieve()` without it, so `project cmd rag-search` raised `TypeError: retrieve() missing 1 required keyword-only argument: 'organization_id'` before any search ran. The unit suite did not catch it - app/commands/rag.py is not in the coverage gate and had no test for this path. The CLI is a tenantless operator path with no acting organization, so it passes `organization_id=None` explicitly, the same as the other tenantless vector-store calls. Reranker resolution with no organization behaves as it did before tenant scoping. Added a regression test asserting search_async scopes retrieval to no organization; it fails (KeyError) against the un-fixed call. Found by codex review of this PR. --- backend/app/commands/rag.py | 1 + backend/tests/test_commands.py | 22 ++++++++++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/backend/app/commands/rag.py b/backend/app/commands/rag.py index f88f17a10..b096c6219 100644 --- a/backend/app/commands/rag.py +++ b/backend/app/commands/rag.py @@ -348,6 +348,7 @@ async def search_async( query=query, collection_name=collection, limit=top_k, + organization_id=None, ) if not results: 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 From 6ae333119b3918fd35d3a8d60761613fb878d168 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Fri, 21 Aug 2026 23:05:28 +0200 Subject: [PATCH 30/37] fix(rag): meter rerank by Cohere's billed search units Rerank spend was derived as ceil(candidate_count / 100), one search unit per 100 candidates. Cohere splits a document past its token threshold into several billable documents, so a request with fewer than 100 candidates but large chunks (ingestion allows chunks up to 8192 chars) consumes several search units while this recorded one - the dashboard and the monthly budget ledger drifting below the provider bill. Read the billed figure from the response instead: meta.billed_units.search_units is what Cohere charged. Every level of that chain is optional, so when it is absent fall back to the old candidate-count estimate. The float is rounded up to whole units and cost stays a Decimal. Tests cover the billed-units path (a small candidate set billed several units), a fractional unit rounding up, and the fallback when the response omits the figure. reranker.py stays at 100% coverage. Found by codex review of this PR. --- backend/app/services/rag/reranker.py | 31 ++++++++++++++++++++++------ backend/tests/test_reranker.py | 31 ++++++++++++++++++++++++++-- docs/file-processing.md | 5 ++++- 3 files changed, 58 insertions(+), 9 deletions(-) diff --git a/backend/app/services/rag/reranker.py b/backend/app/services/rag/reranker.py index cc935745c..0e7469f2b 100644 --- a/backend/app/services/rag/reranker.py +++ b/backend/app/services/rag/reranker.py @@ -40,8 +40,9 @@ logger = logging.getLogger(__name__) # Cohere bills reranking per "search unit": one query with up to this many -# documents. A request with more is split and billed as several - a query with -# 250 documents is three units. +# 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. @@ -121,7 +122,7 @@ async def rerank( # 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(len(results))) + book_ambient_spend(self._spend_entry(response, len(results))) return [ SearchResult( @@ -149,16 +150,34 @@ async def _release(self, client: AsyncClientV2) -> None: except Exception: logger.warning("[RERANK] Closing the Cohere client failed", exc_info=True) - def _spend_entry(self, document_count: int) -> SpendEntry: - units = ceil(document_count / _DOCS_PER_SEARCH_UNIT) + 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 * units, + 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) -> BaseReranker | None: """Bind a collection's resolved rerank credential to a concrete reranker. diff --git a/backend/tests/test_reranker.py b/backend/tests/test_reranker.py index dc417a8e3..88b4e0a02 100644 --- a/backend/tests/test_reranker.py +++ b/backend/tests/test_reranker.py @@ -26,11 +26,18 @@ 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]) -> AsyncMock: - """A fake Cohere client whose rerank returns these (index, score) items.""" +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 @@ -84,6 +91,7 @@ async def test_a_rerank_books_a_priced_nonzero_cost_to_the_active_ledger(self): 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() @@ -92,6 +100,25 @@ async def test_more_than_one_hundred_documents_bills_more_than_one_search_unit(s 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")) diff --git a/docs/file-processing.md b/docs/file-processing.md index a1acf9490..079dd5f9e 100644 --- a/docs/file-processing.md +++ b/docs/file-processing.md @@ -500,7 +500,10 @@ search books automatically — but `SpendLedger.record()` prices through 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) and booked with `book_ambient_spend`, landing +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 From 1ad0015863a8f6dec8d0cc867a71a9a3cc5b7f5a Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Fri, 21 Aug 2026 23:12:18 +0200 Subject: [PATCH 31/37] fix(vault): report knowledge-base bindings in secret usage A secret's usage listing ("used_by" / what breaks if I delete this) was built solely from agents_using(), which searches agent draft specs. A Cohere key used only as a knowledge base's rerank credential - or an OpenAI key used only as a KB embedding credential - therefore read as "not used yet", inviting an admin to delete a key a collection was actively resolving. Both references are SET NULL on delete, so the deletion does not error; it silently stops reranking or embedding for every bound collection. Add knowledge_bases_using(organization_id, secret_id) covering both embedding_secret_id and rerank_secret_id, and fold its rows into used_by as kind="knowledge_base" beside the agents. The embedding gap predates this branch - KB secret references were never surfaced to the vault - so the one query closes both rather than only the rerank case this PR added. - schemas/secret.py: SecretUsage.kind widens to "agent" | "knowledge_base"; the shape is unchanged, so the frontend contract is additive (the table already renders usage names without branching on kind). - frontend types mirror the widened union. - docs/secrets.md documents KB embedding/rerank bindings and that the usage listing reports them. Tests: a key bound only by a KB is reported (unit), and knowledge_bases_using finds both embedding and rerank bindings scoped to the organization (integration). Found by codex review of this PR. --- backend/app/repositories/knowledge_base.py | 27 ++++++++++ backend/app/schemas/secret.py | 2 +- backend/app/services/organization_secret.py | 14 ++++- .../test_collection_name_tenant_isolation.py | 42 +++++++++++++++ backend/tests/test_secrets.py | 52 +++++++++++++++++++ docs/secrets.md | 8 +++ frontend/src/types/secrets.ts | 2 +- 7 files changed, 144 insertions(+), 3 deletions(-) diff --git a/backend/app/repositories/knowledge_base.py b/backend/app/repositories/knowledge_base.py index 98497bc20..7d48340fe 100644 --- a/backend/app/repositories/knowledge_base.py +++ b/backend/app/repositories/knowledge_base.py @@ -188,6 +188,33 @@ async def list_by_collection_name(db: AsyncSession, collection_name: str) -> lis 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: diff --git a/backend/app/schemas/secret.py b/backend/app/schemas/secret.py index 67beac4e7..fb292822c 100644 --- a/backend/app/schemas/secret.py +++ b/backend/app/schemas/secret.py @@ -72,7 +72,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/services/organization_secret.py b/backend/app/services/organization_secret.py index af0f70ab2..987ab7864 100644 --- a/backend/app/services/organization_secret.py +++ b/backend/app/services/organization_secret.py @@ -29,7 +29,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 @@ -126,6 +131,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, @@ -143,6 +151,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/tests/integration/test_collection_name_tenant_isolation.py b/backend/tests/integration/test_collection_name_tenant_isolation.py index cfeb9d0fc..6d5cc6eb8 100644 --- a/backend/tests/integration/test_collection_name_tenant_isolation.py +++ b/backend/tests/integration/test_collection_name_tenant_isolation.py @@ -117,3 +117,45 @@ async def test_an_organization_without_a_row_for_the_name_resolves_nothing(db) - 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_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/test_secrets.py b/backend/tests/test_secrets.py index 31b53335c..24d78f329 100644 --- a/backend/tests/test_secrets.py +++ b/backend/tests/test_secrets.py @@ -31,6 +31,7 @@ ) from app.core.vault import VaultScope 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 @@ -269,12 +270,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/docs/secrets.md b/docs/secrets.md index 5a21cd68e..bf92ef03a 100644 --- a/docs/secrets.md +++ b/docs/secrets.md @@ -110,6 +110,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/src/types/secrets.ts b/frontend/src/types/secrets.ts index 291f56734..935cee891 100644 --- a/frontend/src/types/secrets.ts +++ b/frontend/src/types/secrets.ts @@ -91,7 +91,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; } From 7c985d446525cfe3add331af6ed13720338c8c1a Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Fri, 21 Aug 2026 23:41:53 +0200 Subject: [PATCH 32/37] fix(rag): resolve a search's collection by the authorized knowledge base Scoping resolution to the acting organization (get_for_collection, own-org wins) narrowed the shared-collection-name exposure but did not close it. The access check and the resolution used different tie-breaks on one name: CollectionAccessService returns the first *readable* row, while get_for_collection returns the caller's own-org row unconditionally. When an app-scoped collection and a restricted org collection of the caller's own organization share a collection_name, a member who may read the app collection but holds no grant on the restricted org one was authorized against the app row, yet resolution returned the org row - unsealing and spending its embedding and rerank keys. An intra-organization access-control bypass reaching a key the caller was never granted. Root cause: knowledge_search computed the authorized rows (readable_all) and then kept only their collection_name strings; retrieval and the resolvers re-looked-up by (collection_name, organization_id). Thread the authorized knowledge base id from the search path down to the resolvers, which read that exact row (get_by_id) instead of re-selecting by name. The id comes from the same readable_all that granted access, so resolution can never land on a row access did not. The parameter is optional and defaults to the previous behaviour, so ingestion and the CLI - which choose the row themselves and have no distinct authorized identity - keep the organization-scoped get_for_collection lookup. The agent-run tool identifies its bound collections by name from the spec and is unchanged: it has no per-user access check to diverge from. Threaded through retrieve / retrieve_multi / _recall / _bm25_search / _reranker_for, PgVectorStore.search / _for_collection, build_reranker, and both embeddings_for_collection / reranker_for_collection; the resolver callable types gain the third argument. Cover the refusal: an integration test creates an app collection and a restricted org collection on one name and asserts that resolving by the authorized app row returns the app config and never the org key, where a name+organization lookup returns the org row. Unit tests assert the id reaches the resolvers (retrieval) and that a given id reads that row and skips the name lookup (both resolvers, at 100% on the gated modules). Refs #913. Found by codex review of this PR. --- backend/app/services/embedding_resolution.py | 21 +++++--- backend/app/services/knowledge_search.py | 10 +++- backend/app/services/rag/reranker.py | 10 +++- backend/app/services/rag/retrieval.py | 51 ++++++++++++++---- backend/app/services/rag/vectorstore.py | 16 ++++-- backend/app/services/rerank_resolution.py | 22 +++++--- backend/app/worker/tasks/rag_tasks.py | 8 ++- .../test_collection_name_tenant_isolation.py | 50 ++++++++++++++++++ backend/tests/test_embedding_resolution.py | 28 ++++++++++ backend/tests/test_ingestion_embedding_key.py | 8 ++- backend/tests/test_rerank_resolution.py | 29 +++++++++++ backend/tests/test_retrieval_reranking.py | 52 ++++++++++++++++++- docs/file-processing.md | 10 ++++ 13 files changed, 280 insertions(+), 35 deletions(-) diff --git a/backend/app/services/embedding_resolution.py b/backend/app/services/embedding_resolution.py index 6b4b202e6..bdcdacd25 100644 --- a/backend/app/services/embedding_resolution.py +++ b/backend/app/services/embedding_resolution.py @@ -134,7 +134,7 @@ def describe(self, collection_name: str) -> str: async def embeddings_for_collection( - collection_name: str, organization_id: UUID | None + collection_name: str, organization_id: UUID | None, knowledge_base_id: UUID | None = None ) -> ResolvedEmbeddings | None: """Resolve one collection's embedding model and credential, for one organization. @@ -143,14 +143,21 @@ async def embeddings_for_collection( gotten. Opens its own session because the store embeds from places with no request in sight: a worker mid-ingestion, a capability mid-run. - `organization_id` is required and scopes the resolution: `collection_name` - is not unique across tenants, so resolving by name alone could return - and - unseal and bill - another organization's key (#913). The caller passes the - organization the search or ingest is acting for; `None` only where there is - no tenant (a CLI ingest). + `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. """ async with get_db_context() as db: - kb = await knowledge_base_repo.get_for_collection(db, collection_name, organization_id) + 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 api_key, key_source = await _api_key_for(db, kb) diff --git a/backend/app/services/knowledge_search.py b/backend/app/services/knowledge_search.py index c8b0e429d..a83a5ef3e 100644 --- a/backend/app/services/knowledge_search.py +++ b/backend/app/services/knowledge_search.py @@ -66,7 +66,13 @@ async def search(self, ctx: AuthContext, request: RAGSearchRequest) -> list[Sear refused search never reaches a paid call. """ names = request.collection_names or [request.collection_name] - collections = [kb.collection_name for kb in await self.access.readable_all(ctx, names)] + # 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) @@ -81,6 +87,7 @@ async def search(self, ctx: AuthContext, request: RAGSearchRequest) -> list[Sear 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( @@ -90,6 +97,7 @@ async def search(self, ctx: AuthContext, request: RAGSearchRequest) -> list[Sear 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, diff --git a/backend/app/services/rag/reranker.py b/backend/app/services/rag/reranker.py index 0e7469f2b..59afb66ad 100644 --- a/backend/app/services/rag/reranker.py +++ b/backend/app/services/rag/reranker.py @@ -179,7 +179,9 @@ def _billed_units(response: object, candidate_count: int) -> int: return ceil(candidate_count / _DOCS_PER_SEARCH_UNIT) -async def build_reranker(collection_name: str, organization_id: UUID | None) -> BaseReranker | None: +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 @@ -188,8 +190,12 @@ async def build_reranker(collection_name: str, organization_id: UUID | None) -> `/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) + 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 57246a0e7..97fdb3401 100644 --- a/backend/app/services/rag/retrieval.py +++ b/backend/app/services/rag/retrieval.py @@ -19,7 +19,7 @@ # 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], Awaitable[BaseReranker | None]] +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 @@ -46,6 +46,7 @@ async def retrieve( filter: str = "", *, organization_id: UUID | None, + knowledge_base_id: UUID | None = None, ) -> list[SearchResult]: pass @@ -97,7 +98,12 @@ def _rrf_fuse( ] async def _bm25_search( - self, query: str, collection_name: str, limit: int, organization_id: UUID | None + 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: @@ -108,6 +114,7 @@ async def _bm25_search( query=query, limit=min(limit * 10, 100), organization_id=organization_id, + knowledge_base_id=knowledge_base_id, ) if not all_results: return [] @@ -130,19 +137,24 @@ async def _bm25_search( ] async def _reranker_for( - self, collection_name: str, organization_id: UUID | None + 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. - `organization_id` scopes the resolution: a shared collection name must - resolve the caller's own rerank config, never another tenant's (#913). + `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) + return await self._reranker_resolver(collection_name, organization_id, knowledge_base_id) @staticmethod def _shared_reranker(rerankers: list[BaseReranker | None]) -> BaseReranker | None: @@ -192,8 +204,9 @@ async def retrieve( filter: str = "", *, organization_id: UUID | None, + knowledge_base_id: UUID | None = None, ) -> list[SearchResult]: - reranker = await self._reranker_for(collection_name, organization_id) + 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, @@ -203,6 +216,7 @@ async def retrieve( filter, fetch_multiplier=multiplier, organization_id=organization_id, + knowledge_base_id=knowledge_base_id, ) return await self._rank_and_truncate(reranker, query, candidates, limit) @@ -216,6 +230,7 @@ async def _recall( *, fetch_multiplier: int, organization_id: UUID | None, + knowledge_base_id: UUID | None = None, ) -> list[SearchResult]: """Vector (and optionally BM25) recall, filtered and deduplicated. @@ -246,6 +261,7 @@ async def _recall( filter_expr=filter, limit=limit * fetch_multiplier, organization_id=organization_id, + knowledge_base_id=knowledge_base_id, ) search_time = time.time() - start_time @@ -257,7 +273,7 @@ async def _recall( if self._hybrid_enabled: bm25_results = await self._bm25_search( - query, collection_name, limit * fetch_multiplier, organization_id + query, collection_name, limit * fetch_multiplier, organization_id, knowledge_base_id ) if bm25_results: pipeline_results = self._rrf_fuse(pipeline_results, bm25_results) @@ -322,6 +338,7 @@ async def retrieve_multi( 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. @@ -351,12 +368,25 @@ async def retrieve_multi( merge - each collection's top `limit`, fused, sorted, deduplicated, truncated. """ - rerankers = [await self._reranker_for(name, organization_id) for name in collection_names] + # 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: + for name, kb_id in zip(collection_names, kb_ids): recalled = await self._recall( query, name, @@ -365,6 +395,7 @@ async def retrieve_multi( "", fetch_multiplier=multiplier, organization_id=organization_id, + knowledge_base_id=kb_id, ) all_results.extend(recalled if reranker else recalled[:limit]) diff --git a/backend/app/services/rag/vectorstore.py b/backend/app/services/rag/vectorstore.py index 532033585..1cdcb721d 100644 --- a/backend/app/services/rag/vectorstore.py +++ b/backend/app/services/rag/vectorstore.py @@ -57,6 +57,7 @@ async def search( filter_expr: str = "", *, organization_id: UUID | None, + knowledge_base_id: UUID | None = None, ) -> list[SearchResult]: pass @@ -189,7 +190,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, UUID | None], 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 @@ -286,7 +287,7 @@ def _table(self, name: str) -> str: return f"{VECTOR_TABLE_PREFIX}{name}" async def _for_collection( - self, name: str, organization_id: UUID | None + 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. @@ -303,7 +304,7 @@ async def _for_collection( The recorded width wins over the catalog's: the table was created at that number. """ - resolved = await self._resolver(name, organization_id) + 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) @@ -433,6 +434,7 @@ async def search( 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. @@ -442,11 +444,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, organization_id) + 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 diff --git a/backend/app/services/rerank_resolution.py b/backend/app/services/rerank_resolution.py index bd472fcba..d77286e1d 100644 --- a/backend/app/services/rerank_resolution.py +++ b/backend/app/services/rerank_resolution.py @@ -102,7 +102,7 @@ def __repr__(self) -> str: async def reranker_for_collection( - collection_name: str, organization_id: UUID | None + 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. @@ -112,13 +112,23 @@ async def reranker_for_collection( 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. - `organization_id` scopes the resolution: `collection_name` is not unique - across tenants, so resolving by name alone could read another organization's - rerank config and unseal its key (#913). The caller passes the organization - the search is acting for. + `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_for_collection(db, collection_name, organization_id) + 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) diff --git a/backend/app/worker/tasks/rag_tasks.py b/backend/app/worker/tasks/rag_tasks.py index 4c31ec8e2..22e7ac605 100644 --- a/backend/app/worker/tasks/rag_tasks.py +++ b/backend/app/worker/tasks/rag_tasks.py @@ -84,9 +84,13 @@ def _announcing_resolver() -> EmbeddingResolver: announced: set[str] = set() async def resolve( - collection_name: str, organization_id: UUID | None + collection_name: str, + organization_id: UUID | None, + knowledge_base_id: UUID | None = None, ) -> ResolvedEmbeddings | None: - resolved = await embeddings_for_collection(collection_name, organization_id) + resolved = await embeddings_for_collection( + collection_name, organization_id, knowledge_base_id + ) if ( resolved is not None and resolved.key_source.is_degraded diff --git a/backend/tests/integration/test_collection_name_tenant_isolation.py b/backend/tests/integration/test_collection_name_tenant_isolation.py index 6d5cc6eb8..cf4daeb73 100644 --- a/backend/tests/integration/test_collection_name_tenant_isolation.py +++ b/backend/tests/integration/test_collection_name_tenant_isolation.py @@ -133,6 +133,56 @@ def _kb_row(org: Organization, name: str, **secret_ids: uuid.UUID) -> KnowledgeB ) +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 diff --git a/backend/tests/test_embedding_resolution.py b/backend/tests/test_embedding_resolution.py index 2fae4f67e..81cd6773c 100644 --- a/backend/tests/test_embedding_resolution.py +++ b/backend/tests/test_embedding_resolution.py @@ -107,6 +107,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 dde67a409..6f8d0b463 100644 --- a/backend/tests/test_ingestion_embedding_key.py +++ b/backend/tests/test_ingestion_embedding_key.py @@ -225,7 +225,7 @@ async def test_two_collections_on_one_key_do_not_share_each_others_name(self): ), } store._resolver = AsyncMock( - side_effect=lambda name, organization_id=None: resolutions[name] + side_effect=lambda name, organization_id=None, knowledge_base_id=None: resolutions[name] ) origins = [] @@ -368,7 +368,11 @@ 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, organization_id=None: 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, ): diff --git a/backend/tests/test_rerank_resolution.py b/backend/tests/test_rerank_resolution.py index 07009eb10..44bd6b7ea 100644 --- a/backend/tests/test_rerank_resolution.py +++ b/backend/tests/test_rerank_resolution.py @@ -123,6 +123,35 @@ async def test_the_secret_is_only_ever_looked_up_within_the_collections_org(self 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. diff --git a/backend/tests/test_retrieval_reranking.py b/backend/tests/test_retrieval_reranking.py index 7aa0bb725..a15256c31 100644 --- a/backend/tests/test_retrieval_reranking.py +++ b/backend/tests/test_retrieval_reranking.py @@ -10,6 +10,7 @@ from __future__ import annotations from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 import pytest @@ -62,7 +63,9 @@ def _service_by_name(store: MagicMock, mapping: dict[str, BaseReranker | None]) settings = MagicMock() settings.enable_hybrid_search = False - async def resolver(name: str, organization_id: object = None) -> BaseReranker | None: + 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) @@ -117,6 +120,53 @@ async def test_the_collection_stamp_survives_reranking(self): 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. diff --git a/docs/file-processing.md b/docs/file-processing.md index 079dd5f9e..7d4295a11 100644 --- a/docs/file-processing.md +++ b/docs/file-processing.md @@ -457,6 +457,16 @@ 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 on the search path, though: an `app` +collection everyone may read and a restricted `org` collection of the same name +are both the caller's to resolve, but access may have authorized only the first. +So the search path passes the **authorized** knowledge base's id down to the +resolvers, which read that exact row rather than looking one up by name — the id +comes from the same `readable_all` that granted access, so resolution can never +land on a row access did not (#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 From cfdc96e63dc5b9ea6f69786ffc5bea73af683a6d Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Sat, 22 Aug 2026 00:02:58 +0200 Subject: [PATCH 33/37] test(rag): widen the reserved-names fake resolver to the new signature Threading the authorized knowledge base id gave the embedding resolver a third argument, and `PgVectorStore._for_collection` now calls the resolver with three. This integration test's stub resolver still took two, so it raised `TypeError: _no_collection_of_its_own() takes from 1 to 2 positional arguments but 3 were given` - a failure invisible locally (no database) and caught only by the CI test job. Widen the stub to accept the ignored `knowledge_base_id`. Verified against a real pgvector: the whole tests/integration/ suite passes (584), including this file and the app-vs-restricted-org cover-the-refusal test. --- backend/tests/integration/test_vector_store_reserved_names.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/backend/tests/integration/test_vector_store_reserved_names.py b/backend/tests/integration/test_vector_store_reserved_names.py index bd031bcce..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, organization_id: object = None) -> 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 From 2c72d655df0473aa2c805f024262ee0cfd770a42 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Sun, 23 Aug 2026 23:16:52 +0200 Subject: [PATCH 34/37] fix(rag): renumber the rerank migrations onto main's 0055 Main took the triggers and sandbox stacks since the last re-parent, so its chain now runs to 0055_sandbox_operations and the rerank pair - kept at 0046/0047 from the previous renumbering - forked it at 0045 again, this time also colliding with main's own 0046-0047 by number. Same fix as 55bce968 and a5f3a7c6, one merge later: 0046_knowledge_base_rerank becomes 0056 (down 0055_sandbox_operations) and 0045->0047 ingestion_spend_source becomes 0057. Verified against a real pgvector 16: alembic heads reports one head, tests/test_migrations.py passes whole (upgrade, downgrade, cycle, current-matches-head), tests/test_migration_chain.py passes, the rerank suites (51 tests) and the KB-dialog/onboarding frontend suites pass on the merged tree. Refs #911 --- ...ledge_base_rerank.py => 0056_knowledge_base_rerank.py} | 8 ++++---- ...ion_spend_source.py => 0057_ingestion_spend_source.py} | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) rename backend/alembic/versions/{0046_knowledge_base_rerank.py => 0056_knowledge_base_rerank.py} (91%) rename backend/alembic/versions/{0047_ingestion_spend_source.py => 0057_ingestion_spend_source.py} (85%) diff --git a/backend/alembic/versions/0046_knowledge_base_rerank.py b/backend/alembic/versions/0056_knowledge_base_rerank.py similarity index 91% rename from backend/alembic/versions/0046_knowledge_base_rerank.py rename to backend/alembic/versions/0056_knowledge_base_rerank.py index 65bde54f5..f97096462 100644 --- a/backend/alembic/versions/0046_knowledge_base_rerank.py +++ b/backend/alembic/versions/0056_knowledge_base_rerank.py @@ -14,8 +14,8 @@ the embedding key there is no deployment fallback - a reranker with no key is simply off. -Revision ID: 0046_knowledge_base_rerank -Revises: 0045_audit_impersonator +Revision ID: 0056_knowledge_base_rerank +Revises: 0055_sandbox_operations Create Date: 2026-08-18 """ @@ -26,8 +26,8 @@ from alembic import op -revision: str = "0046_knowledge_base_rerank" -down_revision: str | None = "0045_audit_impersonator" +revision: str = "0056_knowledge_base_rerank" +down_revision: str | None = "0055_sandbox_operations" branch_labels: str | Sequence[str] | None = None depends_on: str | Sequence[str] | None = None diff --git a/backend/alembic/versions/0047_ingestion_spend_source.py b/backend/alembic/versions/0057_ingestion_spend_source.py similarity index 85% rename from backend/alembic/versions/0047_ingestion_spend_source.py rename to backend/alembic/versions/0057_ingestion_spend_source.py index 66f71be27..f39ef95d2 100644 --- a/backend/alembic/versions/0047_ingestion_spend_source.py +++ b/backend/alembic/versions/0057_ingestion_spend_source.py @@ -9,8 +9,8 @@ Every row that predates the column is indexing, so `server_default` backfills them to `'ingestion'` without a data migration. -Revision ID: 0047_ingestion_spend_source -Revises: 0046_knowledge_base_rerank +Revision ID: 0057_ingestion_spend_source +Revises: 0056_knowledge_base_rerank Create Date: 2026-08-20 """ @@ -21,8 +21,8 @@ from alembic import op -revision: str = "0047_ingestion_spend_source" -down_revision: str | None = "0046_knowledge_base_rerank" +revision: str = "0057_ingestion_spend_source" +down_revision: str | None = "0056_knowledge_base_rerank" branch_labels: str | Sequence[str] | None = None depends_on: str | Sequence[str] | None = None From 11c8d47f778ffb579f885234ef4b2cb0da8da77c Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Wed, 26 Aug 2026 19:08:02 +0200 Subject: [PATCH 35/37] fix(rag): carry bound knowledge-base ids through agent retrieval The direct-search fix (#911) threaded the authorized knowledge base id into resolution, but the agent-run path still transported collection names alone. AgentRunnerService._collection_names() resolved spec.collection_ids to names and dropped the ids, and search_knowledge_base supplied only organization_id - so when an agent binds a collection whose collection_name is shared by another row in the same organization, get_for_collection() could select the other row and apply its embedding/rerank configuration, unsealing and billing a key the agent publisher was never granted. The same intra-organization exposure the search path closed, reachable through a bound agent instead of a direct search. Carry the bound id beside the name the whole way: - _collection_names becomes _bound_collections, returning (name, id) pairs; the runner puts both into resources as kb_collection_names / kb_collection_ids. - AgentDeps carries kb_collection_ids aligned with kb_collection_names; the factory reads it from resources. - The delegation path carries it too: ResolvedSubagent.collection_ids, and the clone the library hands a delegate has both put back. - The knowledge toolset passes ctx.deps.kb_collection_ids, and search_knowledge_base forwards knowledge_base_id(s) to retrieve/retrieve_multi. Ids are used only when their length matches the resolved names; the nameless _active_kb_collections fallback (which nothing sets) therefore resolves by organization, exactly as before. Tests: the bound id reaches retrieve on the single and multi paths, and a length mismatch drops the ids rather than pinning the wrong one. Verified the whole backend suite plus the 100% gate against a real pgvector. Refs #913. Found by codex review of this PR. --- .../agents/capabilities/knowledge/_search.py | 11 ++++ .../agents/capabilities/knowledge/_toolset.py | 4 ++ .../capabilities/subagents/_capability.py | 1 + backend/app/agents/deps.py | 16 +++-- backend/app/agents/factory.py | 1 + backend/app/agents/subagent_runtime.py | 4 +- backend/app/services/agent_runner.py | 32 +++++++--- backend/tests/test_capability_edges.py | 58 +++++++++++++++++++ docs/file-processing.md | 21 ++++--- 9 files changed, 124 insertions(+), 24 deletions(-) diff --git a/backend/app/agents/capabilities/knowledge/_search.py b/backend/app/agents/capabilities/knowledge/_search.py index 8292d9f57..311eb4b60 100644 --- a/backend/app/agents/capabilities/knowledge/_search.py +++ b/backend/app/agents/capabilities/knowledge/_search.py @@ -103,6 +103,7 @@ async def search_knowledge_base( top_k: int = 5, *, organization_id: UUID | None, + kb_collection_ids: list[UUID] | None = None, ) -> str: """Search the knowledge base and return formatted results. @@ -115,11 +116,19 @@ async def search_knowledge_base( 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: @@ -129,6 +138,7 @@ async def search_knowledge_base( 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( @@ -136,6 +146,7 @@ async def search_knowledge_base( 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 4d84b4612..be5e722fc 100644 --- a/backend/app/agents/capabilities/knowledge/_toolset.py +++ b/backend/app/agents/capabilities/knowledge/_toolset.py @@ -39,6 +39,10 @@ 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). 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/services/agent_runner.py b/backend/app/services/agent_runner.py index 5b24fb22d..f8c1ec633 100644 --- a/backend/app/services/agent_runner.py +++ b/backend/app/services/agent_runner.py @@ -1317,7 +1317,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=[], )() @@ -1409,13 +1414,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: @@ -1427,8 +1435,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 @@ -1605,8 +1613,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), } @@ -2097,6 +2107,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( @@ -2276,6 +2287,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( @@ -2311,8 +2323,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/tests/test_capability_edges.py b/backend/tests/test_capability_edges.py index 588d89b65..0e591f313 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 @@ -186,6 +187,63 @@ async def test_an_unconfigured_deployment_keeps_saying_what_to_configure(self): 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.""" diff --git a/docs/file-processing.md b/docs/file-processing.md index 7d4295a11..7b469b320 100644 --- a/docs/file-processing.md +++ b/docs/file-processing.md @@ -457,15 +457,18 @@ 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 on the search path, though: an `app` -collection everyone may read and a restricted `org` collection of the same name -are both the caller's to resolve, but access may have authorized only the first. -So the search path passes the **authorized** knowledge base's id down to the -resolvers, which read that exact row rather than looking one up by name — the id -comes from the same `readable_all` that granted access, so resolution can never -land on a row access did not (#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. +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 From 2472ff3cf06f636cb1935bcb4415b6247ef89589 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Wed, 26 Aug 2026 23:42:04 +0200 Subject: [PATCH 36/37] style(rag): wrap the reserved-names test resolver ruff-format wants broken --- backend/tests/integration/test_rag_existence_index.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/backend/tests/integration/test_rag_existence_index.py b/backend/tests/integration/test_rag_existence_index.py index 9ee93ccb3..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, _organization_id: object = None, _kb_id: object = None) -> 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 From 69e42a0f109de7354a6ca4b11d0e9ec3c7e44b88 Mon Sep 17 00:00:00 2001 From: OchnikBartek Date: Wed, 26 Aug 2026 23:55:55 +0200 Subject: [PATCH 37/37] fix(rag): reparent rerank migrations onto main's 0059_invite_fk_ondelete Main merged 0059_invite_fk_ondelete off 0058 while this branch's rerank chain also sat on 0058, so the PR's merge-with-main check saw two 0059 heads. Renumber the branch's two migrations onto main's new head: 0060_knowledge_base_rerank and 0061_ingestion_spend_source. Single linear head verified against a real pgvector (test_migration_chain, test_migrations). --- ...ledge_base_rerank.py => 0060_knowledge_base_rerank.py} | 8 ++++---- ...ion_spend_source.py => 0061_ingestion_spend_source.py} | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) rename backend/alembic/versions/{0059_knowledge_base_rerank.py => 0060_knowledge_base_rerank.py} (90%) rename backend/alembic/versions/{0060_ingestion_spend_source.py => 0061_ingestion_spend_source.py} (85%) diff --git a/backend/alembic/versions/0059_knowledge_base_rerank.py b/backend/alembic/versions/0060_knowledge_base_rerank.py similarity index 90% rename from backend/alembic/versions/0059_knowledge_base_rerank.py rename to backend/alembic/versions/0060_knowledge_base_rerank.py index 6573ab1c5..cc4b9c3d8 100644 --- a/backend/alembic/versions/0059_knowledge_base_rerank.py +++ b/backend/alembic/versions/0060_knowledge_base_rerank.py @@ -14,8 +14,8 @@ the embedding key there is no deployment fallback - a reranker with no key is simply off. -Revision ID: 0059_knowledge_base_rerank -Revises: 0058_backfill_rag_lookup_indexes +Revision ID: 0060_knowledge_base_rerank +Revises: 0059_invite_fk_ondelete Create Date: 2026-08-18 """ @@ -26,8 +26,8 @@ from alembic import op -revision: str = "0059_knowledge_base_rerank" -down_revision: str | None = "0058_backfill_rag_lookup_indexes" +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 diff --git a/backend/alembic/versions/0060_ingestion_spend_source.py b/backend/alembic/versions/0061_ingestion_spend_source.py similarity index 85% rename from backend/alembic/versions/0060_ingestion_spend_source.py rename to backend/alembic/versions/0061_ingestion_spend_source.py index ccc55d91f..80e958c22 100644 --- a/backend/alembic/versions/0060_ingestion_spend_source.py +++ b/backend/alembic/versions/0061_ingestion_spend_source.py @@ -9,8 +9,8 @@ Every row that predates the column is indexing, so `server_default` backfills them to `'ingestion'` without a data migration. -Revision ID: 0060_ingestion_spend_source -Revises: 0059_knowledge_base_rerank +Revision ID: 0061_ingestion_spend_source +Revises: 0060_knowledge_base_rerank Create Date: 2026-08-20 """ @@ -21,8 +21,8 @@ from alembic import op -revision: str = "0060_ingestion_spend_source" -down_revision: str | None = "0059_knowledge_base_rerank" +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