103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
"""知识切片与检索（M3）
|
||
|
||
Revision ID: 0006
|
||
Revises: 0005
|
||
Create Date: 2026-10-29
|
||
|
||
建 `chunk` 表：知识正文按段切开，每片一行，带向量数组。
|
||
|
||
两个刻意的取舍：
|
||
|
||
1. **向量存 `real[]`，不依赖 pgvector**。本机拿不到带 pgvector 的实例（Docker 起不来，
|
||
fnsvr-01 明确不动），但检索链路不能一直停在纸面上。`real[]` + 一条 SQL 余弦函数
|
||
在小数据量下够用；将来换上 pgvector 只需改列类型、加 HNSW 索引，应用层接口不变。
|
||
2. **检索范围只能是 `published`**。草稿是"还没被人认可的话"，进检索等于让 AI 把
|
||
猜测当结论引用出去——这条过滤写在查询里，不靠调用方自觉。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from collections.abc import Sequence
|
||
|
||
from alembic import op
|
||
|
||
revision: str = "0006"
|
||
down_revision: str = "0005"
|
||
branch_labels: Sequence[str] | None = None
|
||
depends_on: Sequence[str] | None = None
|
||
|
||
# 余弦相似度。写成 SQL 函数而不是在 Python 里算，是为了让过滤与排序留在数据库里做——
|
||
# 把整表向量捞回应用算相似度，数据一多就会拖垮 API。
|
||
COSINE_SQL = """
|
||
CREATE OR REPLACE FUNCTION app_cosine(a real[], b real[]) RETURNS double precision
|
||
LANGUAGE plpgsql IMMUTABLE PARALLEL SAFE AS $fn$
|
||
DECLARE
|
||
n int;
|
||
dot double precision := 0;
|
||
na double precision := 0;
|
||
nb double precision := 0;
|
||
BEGIN
|
||
n := LEAST(array_length(a, 1), array_length(b, 1));
|
||
IF n IS NULL THEN
|
||
RETURN NULL;
|
||
END IF;
|
||
FOR i IN 1..n LOOP
|
||
dot := dot + a[i] * b[i];
|
||
na := na + a[i] * a[i];
|
||
nb := nb + b[i] * b[i];
|
||
END LOOP;
|
||
IF na = 0 OR nb = 0 THEN
|
||
RETURN 0;
|
||
END IF;
|
||
RETURN dot / sqrt(na * nb);
|
||
END
|
||
$fn$;
|
||
"""
|
||
|
||
CREATE_SQL = """
|
||
CREATE TABLE chunk (
|
||
id uuid PRIMARY KEY DEFAULT gen_random_uuid(),
|
||
workspace_id uuid NOT NULL,
|
||
item_id uuid NOT NULL REFERENCES item(id) ON DELETE CASCADE,
|
||
idx integer NOT NULL,
|
||
heading varchar(200),
|
||
content text NOT NULL,
|
||
char_count integer NOT NULL DEFAULT 0,
|
||
embedding real[],
|
||
embed_model varchar(64),
|
||
version integer NOT NULL DEFAULT 1,
|
||
created_at timestamptz NOT NULL DEFAULT now(),
|
||
updated_at timestamptz NOT NULL DEFAULT now(),
|
||
CONSTRAINT uq_chunk_item_idx UNIQUE (item_id, idx)
|
||
)
|
||
"""
|
||
|
||
# asyncpg 的执行器不接受多语句，一条 `execute` 只能带一句 SQL（0001 踩过同样的坑）
|
||
INDEX_STATEMENTS = (
|
||
"CREATE INDEX ix_chunk_ws_item ON chunk (workspace_id, item_id)",
|
||
"COMMENT ON TABLE chunk IS '知识正文切片，检索的最小单位'",
|
||
)
|
||
|
||
POLICY_SQL = (
|
||
"ALTER TABLE chunk ENABLE ROW LEVEL SECURITY",
|
||
"CREATE POLICY chunk_isolation ON chunk "
|
||
"USING (workspace_id = app_ws_id()) WITH CHECK (workspace_id = app_ws_id())",
|
||
f"GRANT SELECT, INSERT, UPDATE, DELETE ON chunk TO {APP_ROLE}",
|
||
)
|
||
|
||
|
||
def upgrade() -> None:
|
||
op.execute(COSINE_SQL)
|
||
op.execute(CREATE_SQL)
|
||
for statement in INDEX_STATEMENTS:
|
||
op.execute(statement)
|
||
for statement in POLICY_SQL:
|
||
op.execute(statement)
|
||
|
||
|
||
def downgrade() -> None:
|
||
op.execute("DROP POLICY IF EXISTS chunk_isolation ON chunk")
|
||
op.execute("DROP TABLE IF EXISTS chunk")
|
||
# 函数被别的东西引用时先别删；本轮只有切片在用，留一条注释说明为什么还留着。
|
||
# op.execute("DROP FUNCTION IF EXISTS app_cosine(real[], real[])")
|