126 lines
4.9 KiB
Python
126 lines
4.9 KiB
Python
# -*- coding:utf-8 -*-
|
||
"""P0 探测 v2：VDB(Milvus) /v1/query 正确协议 + expr 标量过滤能力。
|
||
|
||
uapi_seed.sql 实锤协议：kNN = POST /v1/query {colname, vector, pagerows, output_fields}；
|
||
delete = POST /v1/delete {colname, pks}；expr 参数名待实测（v1 探测报错实锤存在 expr 概念）。
|
||
"""
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import sys
|
||
import time
|
||
|
||
WORKDIR = "/d/pipeline/pipeline-app"
|
||
os.chdir(WORKDIR)
|
||
sys.path.insert(0, WORKDIR)
|
||
|
||
from appPublic.folderUtils import ProgramPath # noqa: E402
|
||
from appPublic.jsonConfig import getConfig # noqa: E402
|
||
from appPublic.event_dispatcher import EventDispatcher # noqa: E402
|
||
from sqlor.dbpools import DBPools # noqa: E402
|
||
from ahserver.serverenv import ServerEnv # noqa: E402
|
||
|
||
p = ProgramPath()
|
||
config = getConfig(WORKDIR, NS={'workdir': WORKDIR, 'ProgramPath': p})
|
||
DBPools(config.databases)
|
||
se = ServerEnv()
|
||
se.event_dispatcher = EventDispatcher()
|
||
se.get_module_dbname = lambda m: 'pipeline'
|
||
|
||
TEST_COL = "opp_p0_probe2_%d" % int(time.time())
|
||
DIM = 1024
|
||
|
||
|
||
async def _sql(sql, args=None):
|
||
async with DBPools().sqlorContext("pipeline") as sor:
|
||
recs = await sor.sqlExe(sql, args or {})
|
||
return recs or []
|
||
|
||
|
||
def vec(seed):
|
||
v = [0.0] * DIM
|
||
v[seed] = 1.0
|
||
v[seed + 1] = 0.1
|
||
return v
|
||
|
||
|
||
async def main():
|
||
out = {"steps": [], "col": TEST_COL}
|
||
|
||
def step(name, ok, detail=""):
|
||
out["steps"].append({"name": name, "ok": bool(ok), "detail": str(detail)[:600]})
|
||
print("[%s] %s %s" % ("OK " if ok else "FAIL", name, str(detail)[:400]))
|
||
|
||
recs = await _sql("SELECT baseurl FROM upapp WHERE id='rag-vdb'")
|
||
vdb_base = (recs[0].baseurl or "").rstrip("/") if recs else ""
|
||
if not vdb_base:
|
||
step("vdb_base", False, "未配置")
|
||
return
|
||
|
||
import aiohttp
|
||
timeout = aiohttp.ClientTimeout(total=30)
|
||
async with aiohttp.ClientSession(timeout=timeout) as s:
|
||
async def post(path, payload):
|
||
r = await s.post(vdb_base + path, json=payload)
|
||
return r.status, await r.text()
|
||
|
||
st, body = await post("/v1/createcollection", {"colname": TEST_COL, "fields": [
|
||
{"name": "id", "type": "str", "is_primary": True, "max_length": 64},
|
||
{"name": "vector", "type": "fvector", "dim": DIM},
|
||
{"name": "text", "type": "str", "max_length": 65535},
|
||
{"name": "batch_id", "type": "str", "max_length": 64}],
|
||
"description": "P0 probe2", "metric": "COSINE"})
|
||
step("createcollection(+batch_id字段)", st == 200 and "SUCCEEDED" in body, body)
|
||
|
||
st, body = await post("/v1/upsert", {"colname": TEST_COL, "data": [
|
||
{"id": "t1", "vector": vec(0), "text": "外包人员管理系统", "batch_id": "b1"},
|
||
{"id": "t2", "vector": vec(2), "text": "合同管理系统", "batch_id": "b1"},
|
||
{"id": "t3", "vector": vec(4), "text": "知识库问答", "batch_id": "b2"}]})
|
||
step("upsert(3)", st == 200 and "SUCCEEDED" in body, body)
|
||
|
||
# kNN 检索（生产协议：/v1/query + pagerows）
|
||
st, body = await post("/v1/query", {
|
||
"colname": TEST_COL, "vector": vec(0), "pagerows": 2,
|
||
"output_fields": ["id", "text", "batch_id"]})
|
||
ok = st == 200 and "t1" in body
|
||
step("query.kNN(pagerows=2)", ok, body)
|
||
|
||
# 标量过滤 expr（升级路径①的关键：常驻 collection + batch_id filter）
|
||
for key in ("expr", "filter"):
|
||
st, body = await post("/v1/query", {
|
||
"colname": TEST_COL, key: 'batch_id == "b1"',
|
||
"output_fields": ["id", "text"]})
|
||
ok = st == 200 and "t1" in body and "t3" not in body
|
||
step("query.scalar(%s=...)" % key, ok, body)
|
||
|
||
# kNN + expr 组合（聚类工作集免拷贝的终极验证）
|
||
st, body = await post("/v1/query", {
|
||
"colname": TEST_COL, "vector": vec(0), "pagerows": 3,
|
||
"expr": 'batch_id == "b1"', "output_fields": ["id", "batch_id"]})
|
||
ok = st == 200 and "t1" in body and "t3" not in body
|
||
step("query.kNN+expr", ok, body)
|
||
|
||
# 按主键删除（增量缓存维护用）
|
||
st, body = await post("/v1/delete", {"colname": TEST_COL, "pks": ["t3"]})
|
||
step("delete(pks)", st == 200 and "SUCCEEDED" in body, body)
|
||
st, body = await post("/v1/query", {
|
||
"colname": TEST_COL, "vector": vec(4), "pagerows": 2,
|
||
"output_fields": ["id"]})
|
||
step("delete生效验证(t3应不在)", st == 200 and "t3" not in body, body)
|
||
|
||
# 清理
|
||
st, body = await post("/v1/dropcollection", {"colname": TEST_COL})
|
||
step("dropcollection(清理)", st == 200 and "SUCCEEDED" in body, body)
|
||
|
||
print("=== RESULT_JSON ===")
|
||
print(json.dumps(out, ensure_ascii=False))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
asyncio.run(main())
|
||
except Exception:
|
||
import traceback
|
||
traceback.print_exc()
|
||
sys.exit(1)
|