0mandrock1/rag-tenant
GitHub: 0mandrock1/rag-tenant
基于 PostgreSQL 行级安全策略强制实现租户隔离的多租户 RAG 服务，通过数据库层约束和自动化测试套件从根源上杜绝跨租户数据泄露。
Stars: 0 | Forks: 0
# rag-tenant
一个多租户检索增强生成服务，其租户隔离由 Postgres 行级安全策略强制执行，而非依赖应用程序代码。
本仓库旨在支持的核心主张非常明确且可测试：**如果未来的变更删除了安全策略、配置错了角色或遗漏了过滤条件，测试套件将会失败，而不是导致服务发生数据泄露。** 隔离不应该依赖于 Python 代码在每次查询时都能表现正确。在 `rag_tenant/` 中的任何地方都不存在 `WHERE tenant_id = ...`，并且有一个单元测试断言这种情况永远不会发生。
最核心的产出物是 [`tests/test_isolation.py`](tests/test_isolation.py)。
API 只是为了提供辅助支持。
## 目录
- [论点](#the-argument)
- [快速开始](#quick-start)
- [隔离机制是如何强制执行的](#how-isolation-is-enforced)
- [威胁模型](#threat-model)
- [隔离保证及其局限性](#isolation-guarantees-and-their-limits)
- [隔离测试套件](#the-isolation-test-suite)
- [API](#api)
- [检索路由器](#retrieval-router)
- [压力测试](#load-test)
- [仓库结构](#repository-layout)
- [设计说明](#design-notes)
## 论点
构建多租户 RAG 服务的通常做法是在每个查询的 `WHERE` 子句中加上 `tenant_id`，或者让 ORM 通过作用域 session 来实现。
这在出问题之前是有效的。其失败模式在于，在一个包含四百个查询的代码库中，有一个查询忘记了加过滤条件，而导致的症状是一个客户读取了另一个客户的文档。系统不会崩溃，也不会记录任何日志。代码审查是唯一的防线，而代码审查这种防线在逼近截止日期时会大打折扣。
另一种选择是将约束移入数据库中，这样它将适用于应用程序发出的每一条语句，包括那些未经审查的语句以及目前还不存在的语句：
- 限定租户的表包含 `tenant_id`，并且设置了带有 `FORCE` 选项的 `ENABLE ROW LEVEL SECURITY`。
- 每张表都有一个基于 `current_setting('app.tenant_id', true)` 的策略，通过 `USING` 和 `WITH CHECK` 覆盖 `ALL` 命令。
- 应用程序连接时使用的角色既不是超级用户，也不持有 `BYPASSRLS` 权限。这就是关键所在：RLS 仅对超级用户、具有 `BYPASSRLS` 权限的角色，以及未设置 `FORCE` 的表所有者跳过。应用程序角色不属于这三者中的任何一个，因此查询重写器会将策略谓词追加到它发出的每一条语句中。应用程序层面无法选择退出。
- 每个请求都在显式事务内部以事务局部方式设置 `app.tenant_id`，这样连接池中的连接就不会将一个租户的设置带入下一个租户的请求中。
- 如果某个请求以某种方式在未设置该变量的情况下到达数据库，它将读取到 **0 行** 数据，而不是所有数据。
因为应用程序自身不进行任何过滤，所以一旦删除了某个策略，系统不会悄无声息地降级。它会立即导致测试失败，这将在
[下方进行演示](#the-tests-actually-fail)。
## 快速开始
需要安装支持 Compose 的 Docker。该技术栈使用了自己的项目名称和非默认的主机端口 (5442)，因此不会与现有的 Postgres 发生冲突。
```
cp .env.example .env
docker compose up -d --build
curl -s localhost:8080/healthz
# {"status":"ok"}
```
迁移会在 uvicorn 启动之前由 API 容器的 entrypoint 执行。
配置一个租户并生成一个密钥。这两者都是管理员操作，故意对正在运行的服务不可用：
```
docker compose exec api python -m rag_tenant.cli create-tenant acme --name "Acme Corp"
docker compose exec api python -m rag_tenant.cli create-key acme --name demo
# api_key_id: 6b1e...
# key: rt_fLBHpkxD_...
```
索引一篇文档并提出一个问题：
```
export KEY=rt_...
curl -s -X POST localhost:8080/v1/documents \
-H "Authorization: Bearer $KEY" -H 'Content-Type: application/json' \
-d '{"external_id":"runbook","title":"Failover Runbook",
"content":"When the primary database cluster becomes unreachable, the
on-call engineer promotes the standby replica in the secondary region.
Promotion fences the old primary first to avoid a split brain."}'
curl -s -X POST localhost:8080/v1/query \
-H "Authorization: Bearer $KEY" -H 'Content-Type: application/json' \
-d '{"query":"how does the on-call engineer promote a standby replica"}'
```
```
{
"answer": "Answering '...' from 1 retrieved passage(s):\n[1] Failover Runbook (chunk 0): ...",
"generator": "echo",
"routing": {
"query_class": "hybrid",
"reason": "natural-language query with enough context to warrant dense retrieval fused with full-text",
"embed_tokens": 14,
"embed_tokens_saved": 0
},
"citations": [{"index": 1, "document_title": "Failover Runbook", "matched_by": ["lexical", "vector"], "...": "..."}],
"latency_ms": 46
}
```
无需 LLM API 密钥。默认的生成器是确定性和抽取式的；详情请参见 [回答生成](#answer-generation)。
### 运行测试
```
pip install -r requirements-dev.txt
pytest -m "not integration" # unit suite, no database needed
docker compose up -d postgres
pytest -m integration # the isolation suite
```
## 隔离机制是如何强制执行的
### 角色划分
定义在 [`migrations/002_roles.sql`](migrations/002_roles.sql) 中的三个角色：
| 角色 | 可登录 | 用途 | 为什么单独设立 |
| --- | --- | --- | --- |
| `postgres` | 是 | 迁移、租户配置 | 超级用户，因此无条件绕过 RLS。从不处理 API 流量。 |
| `rag_app` | 是 | 每个 API 请求 | `NOSUPERUSER`, `NOBYPASSRLS`。无法逃脱策略。对 `tenants` 或 `api_keys` 没有 `INSERT` 权限，对任何表都没有 `DELETE` 或 `TRUNCATE` 权限。 |
| `rag_auth` | 否 | 拥有 API 密钥查找函数 | 存在的原因是为了让在确定租户之前必须运行的那个查询被限制在一个单独的 `SECURITY DEFINER` 函数中。`rag_app` 不是其成员，无法 `SET ROLE` 切换到它。 |
`rag_app` 的权限被故意设置得很低
([`005_grants.sql`](migrations/005_grants.sql))。RLS 决定 *哪些行*；
授权决定 *哪些动作*。策略无法阻止应用程序删除允许其查看的行，因此直接不授予应用程序 `DELETE` 权限。
不授予 `TRUNCATE` 权限的原因更为尖锐：**`TRUNCATE` 完全不受行级安全策略的限制**，这将使其成为一种只需一条语句就能实现跨租户数据丢失的元操作。
### 策略
每个限定租户的表都会获得相同的谓词，完整写出如下
([`004_rls.sql`](migrations/004_rls.sql)):
```
ALTER TABLE chunks ENABLE ROW LEVEL SECURITY;
ALTER TABLE chunks FORCE ROW LEVEL SECURITY;
CREATE POLICY chunks_tenant_isolation ON chunks
FOR ALL
USING (tenant_id = nullif(current_setting('app.tenant_id', true), '')::uuid)
WITH CHECK (tenant_id = nullif(current_setting('app.tenant_id', true), '')::uuid);
```
有三个细节至关重要：
- **`FORCE`** 移除了表所有者的豁免权。如果没有它，任何以 schema 所有者身份运行的进程都能看到每个租户的行。
- **`current_setting` 中的 `true`** 使得缺少设置时返回 `NULL` 而不是抛出异常，并且 `nullif(..., '')` 也会将空字符串映射为 `NULL`。`tenant_id = NULL` 的求值结果为 `NULL`，由于它不是 `TRUE`，因此该行会被过滤掉。**其失败模式是返回空结果集，绝对不是全表扫描。**
- **`WITH CHECK`** 将相同的谓词应用于写入操作。如果没有写入隔离，读取隔离将允许一个租户在另一个租户的命名空间中植入行。
该谓词在所有五张表上都被重复声明，而不是封装在辅助函数中。使用辅助函数看起来会更好，但它会将该表达式隐藏在 `pg_policies.qual` 中的一个标识符背后，而测试套件需要对该文本进行断言。
### 每次请求的租户绑定
[`rag_tenant/db.py`](rag_tenant/db.py):
```
with pool.connection() as conn:
with conn.transaction():
with conn.cursor() as cur:
cur.execute(
"SELECT set_config('app.tenant_id', %s, true)", (str(tenant_id),)
)
yield cur
```
`set_config` 的第三个参数是 `is_local`。它是 `SET LOCAL` 的函数形式，其优势在于接受绑定参数，因此租户 ID 永远不会以 SQL 文本的形式到达数据库。
事务作用域使得这在连接池下是安全的。普通的 `SET` 会在物理连接返回到连接池后继续保留，而下一个获取该连接的租户（一个完全不同的租户）将会继承它。这是一种真实的、悄无声息的跨租户读取行为，而 `test_tenant_setting_is_transaction_scoped_and_cannot_leak_via_the_pool` 的存在正是为了专门捕获这种情况，其中包括一个对照组，用于证明该测试能够区分 session 作用域和事务作用域。
调用者接收的是一个 cursor 而不是 connection，因此没有可以直接绕过事务的明显句柄。
### 唯一的例外，以及如何对其进行限制
认证过程是循环的：要应用租户策略，你需要先知道租户；而要获知租户，你必须读取 `api_keys`。总得有什么东西在 `app.tenant_id` 被设置之前读取那张表。
与其给应用程序一个通用的逃生舱，不如将这个例外限制在一个函数中
([`006_auth_lookup.sql`](migrations/006_auth_lookup.sql)):
```
CREATE POLICY api_keys_auth_lookup ON api_keys FOR SELECT TO rag_auth USING (true);
CREATE FUNCTION auth.resolve_api_key(p_key_hash text)
RETURNS TABLE (tenant_id uuid, api_key_id uuid)
LANGUAGE sql STABLE SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
SELECT k.tenant_id, k.api_key_id FROM public.api_keys k
WHERE k.key_hash = p_key_hash AND k.revoked_at IS NULL
$$;
```
`SECURITY DEFINER` 使该函数以其所有者 `rag_auth` 的身份运行，而上面的策略正好允许该角色。`rag_auth` 无法登录，也没有被授予给 `rag_app`，因此接触其权限的唯一方法就是调用此函数——该函数接收一个 key hash 并返回一个租户 ID。它不能进行枚举，不能进行过滤，也不能将 hash 投影出来。调用者必须已经拥有了该密钥，因此该函数不会泄露调用者原本不知道的任何信息。`search_path` 被固定了，因为如果 `SECURITY DEFINER` 函数带有由调用者控制的 `search_path`，就会演变成一个权限提升漏洞。
## 威胁模型
**在范围内。** 这些是该设计旨在抵御的故障。
| 威胁 | 控制措施 |
| --- | --- |
| 编写的查询缺少租户过滤器 | RLS 提供谓词；应用程序从不编写过滤器。 |
| 新的 endpoint 忘记限制其读取范围 | 同上。范围限制不是针对每个 endpoint 的职责。 |
| 连接池化的连接携带了前一个租户的上下文 | 在显式事务内使用 `set_config(..., is_local => true)`；连接池在归还连接时也会执行 `RESET ALL`。 |
| 代码路径在未绑定租户的情况下到达数据库 | 策略求值为 `NULL`；返回零行。 |
| 调用者提供了另一个租户的文档、chunk 或租户 ID | 该 ID 根本不可见；谓词与策略进行了 AND 运算，没有任何表达式可以将其移除。 |
| 调用者构造了一个 embedding，意图与另一个租户的内容发生碰撞 | 向量搜索在索引扫描后受 RLS 过滤。即使距离为零指向了外租户的行，依然什么也不会返回。 |
| 调用者写入带有另一个租户 ID 的行 | `WITH CHECK` 会以 `InsufficientPrivilege` 错误拒绝该操作。 |
| 被盗的 API 密钥 | 限定于一个租户且可撤销；密钥仅以 SHA-256 摘要的形式存储。 |
| 被入侵的应用程序进程为自己伪造凭证 | `rag_app` 对 `tenants` 或 `api_keys` 没有 `INSERT` 权限。配置属于带外管理员操作。 |
| 跨租户数据破坏 | 没有 `DELETE` 授权；没有 `TRUNCATE` 授权，这点至关重要，因为 `TRUNCATE` 会忽略 RLS。 |
| `/metrics` 泄露租户活动 | 没有任何 metric 带有租户标签。每个租户的数据位于 `GET /v1/usage`，受密钥保护和 RLS 限制。 |
**明确不在范围内。** 明确声明这一点是因为一个声称能防御一切的威胁模型毫无价值。
- **被完全入侵的应用程序进程。** `rag_app` 被信任会诚实地设置 `app.tenant_id`。如果在 API 进程内部拥有任意代码执行权限的攻击者，就可以将其设置为任何租户 ID。RLS 防御的是 *查询构建中的 bug*（这是一种常见的失败模式），而不是防御已经控制该进程的攻击者。要控制这种情况，需要针对每个租户的凭证，或者通过一个受信任的代理来发出 `SET` 指令，而这两者都不在本仓库的范围内。
- **迁移超级用户被入侵。** 超级用户根据定义会绕过 RLS。保护这些凭证是一个运维问题，不是这个 schema 能解决的问题。
- **侧信道。** 不对时序、索引统计信息、`EXPLAIN` 成本和错误消息差异进行分析。坚定的攻击者或许能够在不读取某一行的情况下推断出该行的 *存在*。
- **拒绝服务。** 每租户配额是一种成本控制手段，而不是防御措施。它限制了一个租户的支出；但它阻止不了分布式洪流攻击。
- **通过共享生成器产生的内容级泄露。** 不适用于默认的 `EchoGenerator`，它是纯粹抽取式的。引入 LLM 后端会带来 prompt injection（提示词注入）和跨请求缓存问题，这不在本仓库的解决范围内。
- **静态加密、密钥轮换、审计日志完整性。** 未实现。
## 隔离保证及其局限性
以精确的术语说明，实际提供的保证是：
1. 对于由 `rag_app` 发出的、针对 `tenants`、`api_keys`、`documents`、`chunks` 或 `usage_events` 的任何 SQL 语句，只有那些 `tenant_id` 等于当前事务的 `app.tenant_id` 的行是可见的，并且只有这样的行才能被写入。这适用于任意 SQL，而不仅仅是本代码库中的查询——测试套件通过手写的恶意查询验证了这一点。
2. 如果 `app.tenant_id` 未设置或为空，则可见的行数为零。
3. 该设置的生命周期不会超过其所在的事务，因此连接的重用不会导致其在不同租户间发生串联。
4. `rag_app` 无法进行权限提升：它不是超级用户，没有 `BYPASSRLS`，不是 `rag_auth` 的成员，对 `public 没有 `CREATE` 权限。
同样精确地说明其局限性：
- 该保证仅针对 **`rag_app` 角色的行可见性**。它不涉及超级用户能看到什么，也不涉及应用程序进程在合法读取行之后会用这些行做什么。
- 它依赖于应用程序绑定了 *正确的* 租户。RLS 强制执行“每个事务对应一个租户”；它无法告诉你这个绑定是否正确。这种联系——从 API 密钥到租户 ID——属于应用程序代码，代码量很小，且位于 [`rag_tenant/auth.py`](rag_tenant/auth.py) 中。
- `api_keys` 有第二个允许 `rag_auth` 的策略。这确实扩大了攻击面，但其范围受 [上一节](#the-one-exception-and-how-it-is-contained) 中推理的限制。
- 向量和全文索引在各个租户之间是共享的。属于其他租户的行在索引扫描期间会被 *访问* 到，然后再被过滤掉。这样做是正确的，但有一个值得说明的后果：在大型共享 HNSW 索引中，小型租户的召回率可能会下降，因为近似搜索返回的 `ef_search` 个候选结果可能大部分都被过滤掉了。按租户分区的索引可以解决召回率问题，并且是在规模扩大时显而易见的下一步；但这并不是为了修复安全问题，因为安全问题已经得到了解决。
- 配额衍生自 `usage_events`，而不是共享的计数器存储，因此它的隔离程度与其他一切事物完全相同——但在并发条件下它也是近似值，因为两个同时发出的请求可能都会读取到低于限制的计数值。
## 隔离测试套件
六个命名的测试，每个都编写为安全失效。有几个测试带有 *对照组*：一个故意破坏的变体，用于断言该测试确实能检测到故障。如果一个测试在系统损坏的情况下也能通过，那它就毫无价值。
```
$ pytest tests/test_isolation.py -m integration -v
============================= test session starts ==============================
collecting ... collected 6 items
tests/test_isolation.py::test_tenant_a_key_cannot_read_tenant_b_chunks_through_any_endpoint PASSED [ 16%]
tests/test_isolation.py::test_direct_sql_without_set_local_returns_zero_rows PASSED [ 33%]
tests/test_isolation.py::test_tenant_setting_is_transaction_scoped_and_cannot_leak_via_the_pool PASSED [ 50%]
tests/test_isolation.py::test_crafted_filters_and_embeddings_cannot_reach_another_tenant PASSED [ 66%]
tests/test_isolation.py::test_policies_exist_are_forced_and_the_app_role_cannot_bypass_them PASSED [ 83%]
tests/test_isolation.py::test_property_no_cross_tenant_read_for_random_tenants_and_documents PASSED [100%]
============================== 6 passed in 3.25s ===============================
```
每个测试的作用：
1. **`test_tenant_a_key_cannot_read_tenant_b_chunks_through_any_endpoint`** — 租户 B 摄取了一个带有特殊标记的文档；租户 A 通过 `/v1/query` 查询该标记，并检查了返回的回答和每一个 citation，然后检查 `/v1/usage` 报告的流量是否仅属于自己。它还确认了 A 可以重用 B 的 `external_id`，因为文档身份是按租户划分作用域的。
2. **`test_direct_sql_without_set_local_returns_zero_rows`** — 以 `rag_app` 身份打开一个原始连接，在未绑定租户的情况下对全部五张表进行计数。每个计数必须为零。然后，对照组在同一连接上绑定一个租户，并断言行确实是存在的，这样测试就不会仅仅因为数据库是空的而通过。
3. **`test_tenant_setting_is_transaction_scoped_and_cannot_leak_via_the_pool`**
— 强制连接池大小为 1，以便两个租户在可证明的情况下获得相同的后端 PID，运行租户 A 的事务，然后在下次获取连接时断言 `app.tenant_id` 已消失，且查询不返回任何内容。对照组使用 `is_local => false` 重复该序列，并断言该设置 *确实* 发生了泄露，从而证明该测试能够区分 session 作用域和事务作用域。
4. **`test_crafted_filters_and_embeddings_cannot_reach_another_tenant`** — 包含五次攻击，每次攻击在仅依靠应用层过滤的情况下都会得逞：直接指名道姓地使用受害者的 `tenant_id`；使用 `WHERE true OR 1=1`；使用根据受害者确切 chunk 文本计算出的向量，这是与其内容距离最近的 embedding；通过主键读取 chunk；以及带有受害者 `tenant_id` 标记的 `INSERT`，该操作必须抛出 `InsufficientPrivilege` 异常。
5. **`test_policies_exist_are_forced_and_the_app_role_cannot_bypass_them`** — 对机制而非行为进行断言，这样“隔离测试通过”就永远不会变成“隔离测试通过是因为没有数据”。它检查所有五张表上的 `relrowsecurity` 和 `relforcerowsecurity`，每张表是否都有 `tenant_id` 列，每个策略是否覆盖了 `ALL` 并且带有包含 `current_setting('app.tenant_id'::text, true)` 的 `qual` 和 `with_check`，检查 `rag_app` 既不是超级用户，没有 `BYPASSRLS`，也不是 `rag_auth` 的成员，以及它对任何租户表都不持有 `DELETE` 或 `TRUNCATE` 权限。
6. **`test_property_no_cross_tenant_read_for_random_tenants_and_documents`** — 四个租户，每个租户三份文档，随机化的标记和正文。然后每个租户通过 API 和原始 SQL（`DISTINCT tenant_id`，针对外部标记的 `LIKE ANY` 扫描，以及被断言为不会跨越租户边界的 join）搜索不属于它的标记。失败时会打印出 seed，并可通过 `ISOLATION_SEED` 进行固定。
### 测试确实会失败
通过对实时 schema 进行变异并重新运行套件来验证：
| 变异 | 结果 |
| --- | --- |
| `DROP POLICY chunks_tenant_isolation ON chunks` | **6 个中 6 个失败** |
| `ALTER ROLE rag_app BYPASSRLS` | **6 个中 6 个失败** |
| `ALTER TABLE chunks NO FORCE ROW LEVEL SECURITY` | **6 个中 1 个失败**，并带有明确的消息 `chunks has RLS enabled but not FORCE'd; the table owner would bypass every policy on it` |
删除策略会导致每个测试都失败，因为启用了 RLS 但没有策略的表默认会被拒绝——即拒绝方向。授予 `BYPASSRLS` 会在 *泄露* 方向（这也是至关重要的一点）导致每个测试失败。
移除 `FORCE` 是一种有趣的情况，也是测试 5 存在的原因。只有一个测试注意到了它，因为 `rag_app` 不是 `chunks` 的所有者：它的行为保持不变，所以所有的行为测试仍然通过。这个漏洞是为任何 *以所有者身份* 访问数据库的事物敞开的——比如一个迁移、一个维护脚本、一个重用管理员凭证的未来服务——而从应用程序的角度进行再多的黑盒测试也无法揭示这一点。这正是一个测试对机制而非行为进行断言的确切原因。
单元测试套件添加了不需要数据库的静态防护：迁移是否连续编号，每个租户表是否都被 `FORCE` 并且拥有策略，谓词是否出现在每个 `CREATE POLICY` 语句的 `USING` 和 `WITH CHECK` 中，没有 `GRANT` 授予 `DELETE` 或 `TRUNCATE`，以及 `rag_tenant/` 中没有任何可执行的字符串字面量包含 `WHERE ... tenant_id =` 过滤器。最后一个检查使用 `ast` 解析每个模块并跳过 docstring，因此解释该规则的文字本身不会触发误报。
## API
每个 `/v1` 路由上的身份验证方式均为 `Authorization: Bearer `。密钥经过 SHA-256 hash 处理，并通过 `auth.resolve_api_key` 解析为租户。未知和已撤销的密钥都会返回完全相同的 `401`，因此该 endpoint 不能被用来探测过去存在过哪些密钥。
| 路由 | 用途 |
| --- | --- |
| `POST /v1/documents` | 摄取：分块、向量化、索引。成功返回 `201`，如果 `external_id` 存在但内容不同则返回 `409`，如果是完全相同的内容重放则返回 `200` 及 `already_indexed: true`。 |
| `POST /v1/query` | 路由、检索、生成。返回回答、路由决策及其依据，以及带有匹配检索分支标签的 citation。 |
| `GET /v1/usage` | 特定时间窗口内按租户统计的使用情况：按路由器分类的请求数和 embedding token 数、路由器节省的 token 数，以及当前的配额消耗量。 |
| `GET /healthz` | 存活探测，外加一次 *以 `rag_app` 身份* 进行的真实往返操作，因此角色配置错误将表现为不健康状态，而不是等到第一个请求时才暴露。 |
| `GET /metrics` | Prometheus 格式展示。根据设计，不包含任何租户标签。 |
Handler 是同步的 `def`，因此 Starlette 会在线程池中运行它们，这与底层的同步 psycopg 连接池相匹配。有一项规则在整个过程中起着关键的支撑作用：**在进行 CPU 密集型工作时，不持有任何数据库连接。** 在这里，Embedding 耗时远远超过任何查询，因此每个 Handler 都会在一个极短的事务中进行授权，释放连接，进行 Embedding 操作，然后在第二个事务中进行数据持久化。
### 存储
`tenants`, `api_keys`, `documents`, `chunks`, `usage_events`。Chunks 同时带有一个使用 HNSW 索引（基于余弦距离）的 `vector(384)` embedding，以及一个使用 GIN 索引的自动生成的 `tsvector` 列。Embedding 来自于通过 fastembed 运行的 `BAAI/bge-small-en-v1.5`，该库在 CPU 上运行量化的 ONNX。依赖树中没有 torch，也不应该有：它体积超过 1 GB，但在一台 2 vCPU 的机器上，它买不到任何 130 MB ONNX 模型已经提供不了的东西。
### 回答生成
一个包含两种实现的 `Generator` 协议。`EchoGenerator` 是默认的：确定性的、离线的、纯粹抽取式的——它会陈述检索到的内容并附带 citation 进行引用。这使得隔离套件能够对确切的答案文本进行断言，因为该答案是检索行的一个函数，仅此而已。只有在设置了 `ANTHROPIC_API_KEY` 时才会构造 `AnthropicGenerator`，因此绝对不可能出现因为缺少密钥而在请求期间发生运行时错误的代码路径。
## 检索路由器
大多数 RAG 服务会对每一个传入的查询进行 Embedding。这是一种会在账单上体现出来的浪费：`hi`、`what can you do?` 和 `ERR_CONN_REFUSED` 根本不需要进行 384 维的最近邻搜索，而且其中两个甚至完全不需要检索。
[`rag_tenant/router.py`](rag_tenant/router.py) 使用确定性的字符串检查将每个查询分为三类之一——不需要模型，不需要网络，也没有增加延迟。使用 LLM 进行分类的成本，比它所能节省的 Embedding 成本还要高。
| 分类 | 触发条件 | 原理 |
| --- | --- | --- |
| `no_retrieval` | 问候和确认（允许旁边带有停用词）、有关助手的问题、仅包含停用词的输入 | 语料库中没有任何内容可以回答“hello there”。不触碰索引，不生成 Embedding。 |
| `lexical_only` | 带引号的短语；标识符、错误代码、文件名、版本号或 commit SHA；或者少于三个实际内容词 | 在这里，全文搜索不仅成本更低，**而且效果更好**。Embedding 会模糊确切的 token，而这恰恰与搜索 `ERR_CONN_REFUSED` 的人所期望的完全相反。 |
| `hybrid` | 其他所有情况 | 进行 Embedding，执行向量与全文搜索，通过倒数秩融合（reciprocal rank fusion，`k = 60`）进行融合。 |
基于 rank 而非 score 进行融合规避了一个校准问题：`ts_rank_cd` 是无界的且依赖于语料库，而余弦相似度在 `[-1, 1]` 之间，因此这两者作为数字是无法比较的。
启发式方法会发生误分类，而其失败模式在设计上被限制在单侧：被错误地发送到 `lexical_only` 的查询会返回较差的结果，而错误地发送到 `hybrid` 的查询只会消耗几分之一美分。这两种情况都不能返回另一个租户的行，因为路由发生在数据库之上，而数据库并不关心是哪个分支发出的请求。
### 配额
按租户、按滚动时间窗口计算（默认每 60 秒 120 个请求和 20,000 个 Embedding token，可按租户覆盖）。超过其中任何一个限制都会返回 `429`，并带有根据窗口内最早事件过期时间计算出的 `Retry-After` 标头。
Embedding token 的检查是针对 *预计* 消耗量进行的，因此请求会在模型运行之前而不是运行之后被拒绝。
计数器衍生自 `usage_events` 而非 Redis，按重要性由高到低排列，原因有三：该表与所有其他内容受相同的 RLS 策略限制，因此租户 A 在物理上无法看到租户 B 的流量——而使用共享的 Redis 键空间则会通过键前缀将租户分离重新推回应用程序代码中；这是一个可以少运行的服务；并且导致返回 `429` 的底层数字与 `GET /v1/usage` 报告的数字是一致的，不会产生偏差。
被拒绝的请求会被记录下来，但它们本身不消耗配额，因此不断重试的客户端不会将其自身的恢复节点推得更远。
每一个决策都会落入 `usage_events`，包括 *未* 消耗的 token因此这种节省是被实际测量出来的，而不是口头断言的：
在一次下文提到的压力测试运行之后，查看 `GET /v1/usage`：
```
{
"window_seconds": 3600,
"requests": 1724,
"embed_tokens": 37431,
"embed_tokens_saved": 2696,
"embed_tokens_avoided_pct": 6.72,
"by_class": [
{"query_class": "hybrid", "requests": 837, "embed_tokens": 13145, "embed_tokens_saved": 0},
{"query_class": "ingest", "requests": 24, "embed_tokens": 24286, "embed_tokens_saved": 0},
{"query_class": "lexical_only", "requests": 523, "embed_tokens": 0, "embed_tokens_saved": 1852},
{"query_class": "no_retrieval", "requests": 340, "embed_tokens": 0, "embed_tokens_saved": 844}
]
}
```
如果诚实地解读，这比大标题百分比所暗示的胜利要小，而且方式很有趣。**1,700 个查询中有 863 个（51%）完全跳过了 Embedding 调用**，但它们只占 13,145 个实际花费在查询上的 token 中的 2,696 个：即查询 Embedding token 的 17%。被跳过的查询都很短，因此仅计算 token 数会低估其带来的好处；延迟上的节省才是更大的奖品，在空闲状态下，`no_retrieval` 的响应速度大约比混合响应快 2.5 倍。
总体数据还要低（6.7%），因为摄取操作主导了 token 数量，而摄取永远不会被路由。路由器是一种查询路径优化手段，对语料库本身不起作用。如果报告总体数据而不加这种警告，那只会产生一种让仪表盘看起来很漂亮、但账单依然保持原样的数字。
## 压力测试
`POST /v1/query`，20 个并发用户，60 秒，`wait_time` 为 0.1–0.5s，任务混合比例为 5 个 `hybrid` : 3 个 `lexical_only` : 2 个 `no_retrieval`。每个请求还会断言路由器是否返回了配置文件所预期的分类，因此如果某次运行悄悄地将所有请求都重新分类到低成本的路径中，测试将会失败，而不是报告一个虚假好看的 p95。语料库：由 [`loadtest/seed.py`](loadtest/seed.py) 植入的 24 份文档，`EchoGenerator`（无 LLM 延迟），CPU 上运行 `BAAI/bge-small-en-v1.5`，两个 uvicorn worker。
API 容器被 Compose 的 `deploy.resources.limits` 限制为 **2.0 个 CPU 和 1800 MB**，与目标部署环境相匹配。Postgres 运行在一个同胞容器中，内存上限为 768 MB，但 *未对* CPU 进行限制，这对于部署来说是正确的——数据库通常位于独立的主机上。可以通过 `make loadtest` 进行复现。
**单客户端，无并发**（各 20 个请求）：
| 分类 | p50 | p95 |
| --- | --- | --- |
| `hybrid` | 42 ms | 47 ms |
| `lexical_only` | 19 ms | 22 ms |
| `no_retrieval` | 15 ms | 18 ms |
`hybrid` 和 `lexical_only` 之间约 23 ms 的差距就是查询 Embedding 的时间。这正是路由器所要挽回的成本。
**20 个并发用户：**
| 分类 | p50 | p95 | p99 | 请求数 |
| --- | --- | --- | --- | --- |
| `hybrid` | 200 ms | 420 ms | 510 ms | 1220 |
| `lexical_only` | 140 ms | 350 ms | 420 ms | 733 |
| `no_retrieval` | 130 ms | 310 ms | 380 ms | 439 |
| **汇总** | **160 ms** | **390 ms** | 480 ms | 2392 |
持续吞吐量为 40.1 req/s，共 2392 个请求，0 次失败。
该服务受限于队列而非受限于操作处理能力：20 个并发用户对战 2 个核心是一个很深的队列，在空闲 p50 和负载 p50 之间大约 4 倍的差距正是那个队列导致的，而不是模型。路由器的优势在两列数据中都得以保持。
**更严格的变体。** 如果在让 Postgres 使用目标机器上根本不存在的 CPU 核心的情况下报告“2 vCPU”，未免过于宽松了。[`docker-compose.loadtest.yml`](docker-compose.loadtest.yml) 将 2.0 CPU 的预算在两个容器之间进行了分配（API 为 1.5，Postgres 为 0.5）：
| | p50 | p95 | 吞吐量 |
| --- | --- | --- | --- |
| API 2.0 CPU，Postgres 不受限 | 160 ms | 390 ms | 40.1 req/s |
| API 1.5 + Postgres 0.5（共计 2.0） | 320–410 ms | 680–780 ms | 27–30 req/s |
两种配置均为：零失败。正如预期的那样，在两者之间，空闲时的延迟没有变化——单个请求永远不会使任何一种分配方式达到饱和。
更严格的那行数据以范围的形式给出，因为它是多次重复测试中变动最大的数据。将 Postgres 限制在半个核心内，会使运行过程对主机上正在运行的其他任务变得敏感，多次重复测试的结果落在 27.0 到 30.4 req/s 之间。相比之下，主要行的复现结果非常紧凑：独立的重复测试分别给出了 40.1 和 40.2 req/s，p50 为 160 ms 和 150 ms。
**注意事项，这比数字本身更重要。** 测试主机是一台 16 核工作站，通过 cgroup 对容器进行了限制，而不是一台真正的 2 vCPU VPS；对大型机器进行限制和运行小型机器，在内存带宽或缓存行为上并不是一回事。负载生成器运行在不受限的同一台主机上，并且在某些重复测试期间存在一些轻微的无关负载，这也是更严格的那行数据以范围而不是以定点值引用的主要原因。请将这些数据视为服务在并发下的表现形态——受限于队列，且路由器的节省在两列数据中均可见——而不是将其视为一种规范。
### 一个值得记录的发现
这个基准测试最初的运行速度比上面的数字慢了三到四倍，而原因并不在于服务本身。问题出在 onnxruntime 上：**它是根据它能看到的 CPU 数量来设置其线程池大小的，在容器内部，这是宿主机的核心数，而不是 cgroup 配额。** 两个 uvicorn worker 各自针对 2 个 CPU 的配额生成了一个宿主机宽度的内部操作池，把它们的预算都浪费在了上下文切换上。
该症状随着宿主机负载情况的不同而变化，这正是最初让人感到困惑的原因。在繁忙的主机上，它表现为延迟：在 10.2 req/s 的吞吐量下，汇总的 p50 达到了 1400 ms，p95 达到了 3300 ms，但没有发生失败，日志里也没有任何记录。而在空闲的主机上，同样的错误配置表现为连接被拒绝和重置，此时吞吐量为 16.4 req/s。无论是哪种情况，都没有发生崩溃、重启或 OOM kill。
在 ONNX session 上固定 `threads=1`，并在镜像中设置 `OMP_NUM_THREADS=1` 及其对应的 BLAS 环境变量，同时使 worker 数量与 CPU 配额相匹配后：
| | 汇总 p50 | 汇总 p95 | 吞吐量 |
| --- | --- | --- | --- |
| 默认线程 | 1400 ms | 3300 ms | 10.2–16.4 req/s |
| 固定线程数 | 160 ms | 390 ms | 40.1 req/s |
在容器内，`os.cpu_count()` 依然报告为 16。任何依赖它来设置池大小的程序——无论是 onnxruntime、OpenBLAS，还是线程池执行器——都会犯同样的错误，而对 CPU 进行限制正是让这个 bug 暴露出来的原因。
## 仓库结构
```
migrations/ Numbered SQL. 004_rls.sql is the security boundary.
rag_tenant/
api.py FastAPI routes
db.py Pool and tenant_transaction
auth.py Bearer key to tenant
router.py Retrieval cost gate
quota.py Per-tenant rate and token limits
retrieval.py Lexical, vector, RRF fusion. No tenant filtering, on purpose.
embeddings.py fastembed (ONNX) and a deterministic offline backend
generator.py Generator protocol, EchoGenerator, AnthropicGenerator
migrate.py Migration runner
cli.py Operator commands, superuser only
tests/
test_isolation.py The headline suite
loadtest/ Locust profile
```
CI 会对 `pgvector/pgvector:pg16` 服务容器运行 lint 检查、单元套件以及完整的隔离套件。
## 设计说明
那些有不止一种合理答案的决策。我们采用了更简单的选项，并将其记录下来。
- **迁移以超级用户身份运行；而应用程序从不这样做。** 超级用户会绕过 RLS，这正是为什么迁移和配置操作要与请求处理分离开来，而不是共享同一个连接字符串的原因。
- **SQL 中定义角色权限，角色密码取自环境变量。** 迁移 002 声明了 `rag_app` 可以做什么；运行器随后使用 psycopg 的 SQL 构建器执行 `ALTER ROLE ... PASSWORD`。没有提交任何密钥。
- **没有 ORM。** 为你自动添加 `WHERE tenant_id = ?` 的 ORM 正是本仓库所反对的模式。纯粹的 SQL 让租户过滤器的缺失保持可见且可测试。
- **重新摄取修改后的文档会返回 `409` 而不是替换它。** 替换意味着要删除 chunks，而 `rag_app` 没有 `DELETE` 权限。保持无此权限比支持原地更新更有价值；请使用一个新的 `external_id` 进行摄取。
- **隔离套件使用 `DELETE FROM tenants` 清除数据，而不是 `TRUNCATE`。** 所有内容都从 `tenants` 层层级联，并且 `DELETE` 采用行锁，而 `TRUNCATE` 采用的 `AccessExclusiveLock` 会在连接池连接仍在读取时发生死锁。
- **Token 计数估算为每四个字符一个 token。** 用于配额核算和“节省的 token”数据，这两者都需要一个在各 Embedding 后端之间保持稳定而不是绝对精确的数字。它还使得配额行为在 CI 环境中保持一致，因为 CI 使用的是离线后端。
- **CI 使用确定性哈希 Embedding 模型。** 隔离套件测试的是行可见性，绝不能因为模型主机响应慢而导致失败。它还使得测试 4 中构造 Embedding 的攻击更加尖锐：如果攻击者知道了映射关系，就可以准确地将向量对准受害者的内容。
- **API 密钥使用 SHA-256，而不是 bcrypt 或 argon2。** 那些算法是为人类选择的密码设计的，这类密码的搜索空间很小，以至于工作量因素本身就是防线。而这些密钥是 32 字节的 `secrets.token_urlsafe`；没有任何东西可以暴力破解，如果使用慢速 KDF，只会给每个请求徒增延迟，毫无益处。
- **配额源自 `usage_events`，而不是 Redis。** 推理过程见 [配额](#quota)。为了将租户分离保持在数据库中，接受了并发情况下的近似执行。
- **Prometheus 指标上没有租户标签。** `/metrics` 是未经身份验证的；租户标签会将抓取 endpoint 变成一个侧信道，暴露存在哪些租户以及它们的繁忙程度。
- **基于带有重叠的空白字符进行 Chunking。** Chunk 的质量不是本仓库关注的核心，更智能的分词器会增加一个依赖项和一个故障点，却不会改变隔离机制。
标签：PostgreSQL, Python, RAG, 数据隔离, 无后门, 测试用例, 版权保护, 秘密管理, 行级安全, 请求拦截, 逆向工具