56 lines
3.5 KiB
Python
56 lines
3.5 KiB
Python
"""Real extension smoke test in a new disposable schema of the test database.
|
|
|
|
The 3D vectors are synthetic fixtures, not model embeddings or a recall test.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
from uuid import uuid4
|
|
|
|
import psycopg
|
|
from psycopg import sql
|
|
|
|
|
|
def main() -> None:
|
|
url = os.environ.get("ASSISTANT_TEST_DATABASE_URL", "")
|
|
if not url:
|
|
raise SystemExit("Run pnpm db:up to configure the isolated PostgreSQL test database.")
|
|
schema = "assistant_search_smoke_" + uuid4().hex
|
|
with psycopg.connect(url, connect_timeout=5) as connection:
|
|
major = int(connection.execute("SHOW server_version_num").fetchone()[0]) // 10000
|
|
extensions = dict(connection.execute("SELECT extname,extversion FROM pg_extension WHERE extname IN ('pg_search','vector')"))
|
|
assert major == 17 and extensions == {"pg_search": "0.26.1", "vector": "0.8.6"}, "PostgreSQL extension versions differ from the pinned baseline."
|
|
connection.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema)))
|
|
connection.commit()
|
|
try:
|
|
connection.execute(sql.SQL("SET search_path TO {},public,pg_catalog").format(sql.Identifier(schema)))
|
|
connection.execute("CREATE TABLE documents(id bigint PRIMARY KEY, platform text NOT NULL, body text NOT NULL, embedding vector(3) NOT NULL)")
|
|
for row in [(1, 'xiaohongshu', '建材采购 搜索笔记 查看详情 读取客户评论', '[1,0,0]'),
|
|
(2, 'douyin', '建材装修 视频作者 主页读取', '[0.6,0.8,0]'),
|
|
(3, 'xiaohongshu', '蛋糕烘焙 菜谱教学', '[0,0,1]')]:
|
|
connection.execute("INSERT INTO documents VALUES(%s,%s,%s,%s::vector)", row)
|
|
connection.execute("CREATE INDEX documents_bm25 ON documents USING paradedb (id,(body::pdb.jieba),(platform::pdb.literal))")
|
|
connection.execute("ANALYZE documents")
|
|
connection.commit()
|
|
lexical = connection.execute("SELECT id,pdb.score(id) FROM documents WHERE body ||| %s AND platform=%s ORDER BY pdb.score(id) DESC,id LIMIT 3", ('建材', 'xiaohongshu')).fetchall()
|
|
assert lexical and [row[0] for row in lexical] == [1] and lexical[0][1] > 0, "Chinese BM25 scope/ranking mismatch."
|
|
plan = connection.execute("EXPLAIN (FORMAT JSON) SELECT id,pdb.score(id) FROM documents WHERE body ||| %s ORDER BY pdb.score(id) DESC LIMIT 3", ('建材',)).fetchone()[0]
|
|
assert 'Custom Scan' in str(plan) and ('ParadeDB' in str(plan) or 'Search' in str(plan)), "BM25 did not use its search index."
|
|
nearest = connection.execute("SELECT id FROM documents WHERE platform=%s ORDER BY embedding <=> %s::vector,id LIMIT 1", ('xiaohongshu', '[1,0,0]')).fetchone()[0]
|
|
assert nearest == 1, "Fixture vector ordering mismatch."
|
|
print(json.dumps({"result": "passed", "postgres_major": major, "extensions": extensions,
|
|
"chinese_bm25": True, "platform_filter": True, "search_index_plan": True,
|
|
"fixture_vector_ordering": True, "real_embedding_service": False}))
|
|
finally:
|
|
connection.rollback()
|
|
connection.execute(sql.SQL("DROP SCHEMA {} CASCADE").format(sql.Identifier(schema)))
|
|
connection.commit()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except psycopg.Error:
|
|
raise SystemExit("PostgreSQL extension check failed; check the local container and connection configuration (driver details withheld).") from None
|