195 lines
8.9 KiB
Python
195 lines
8.9 KiB
Python
# -*- coding:utf-8 -*-
|
||
"""P0 探测：VDB(Milvus) 四件套 + embedding 在线引擎连通性（只读探测 + 测试 collection 自清理）。
|
||
|
||
在 pipeline-app 应用根目录执行：./py3/bin/python pkgs/pipeline-opportunity/scripts/p0_probe_vdb_embed.py
|
||
探测项：
|
||
1. upapp.rag-vdb baseurl / rag_engine_configs embedding 配置（key 不回显）
|
||
2. VDB: createcollection(dim=1024, COSINE) → upsert 3条 → search(kNN top_k) → query(标量 filter 表达式能力)
|
||
3. VDB: dropcollection 是否存在（决定批次 collection 清理策略）
|
||
4. embedding OpenAI 兼容 /embeddings：2条文本 → 维度
|
||
全程使用测试专用 collection 名 opp_p0_probe_<ts>，结束尝试清理；任何一步失败如实打印不掩盖。
|
||
"""
|
||
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_probe_%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 []
|
||
|
||
|
||
async def main():
|
||
out = {"steps": []}
|
||
|
||
def step(name, ok, detail=""):
|
||
out["steps"].append({"name": name, "ok": bool(ok), "detail": str(detail)[:500]})
|
||
print("[%s] %s %s" % ("OK " if ok else "FAIL", name, str(detail)[:300]))
|
||
|
||
# ── 1. 配置读取 ──
|
||
recs = await _sql("SELECT baseurl FROM upapp WHERE id='rag-vdb'")
|
||
vdb_base = (recs[0].baseurl or "").rstrip("/") if recs else ""
|
||
step("vdb_base(upapp.rag-vdb)", bool(vdb_base), vdb_base or "未配置")
|
||
|
||
erecs = await _sql("SELECT engine_type, model_name, endpoint_url, status FROM rag_engine_configs WHERE status='active'")
|
||
emb_cfg = None
|
||
for r in erecs:
|
||
if r.engine_type == "embedding":
|
||
emb_cfg = {"model": r.model_name, "base": (r.endpoint_url or "").rstrip("/")}
|
||
step("embedding_cfg(rag_engine_configs)", bool(emb_cfg), json.dumps(emb_cfg, ensure_ascii=False) if emb_cfg else "无 active embedding 配置")
|
||
|
||
import aiohttp
|
||
timeout = aiohttp.ClientTimeout(total=30)
|
||
|
||
# ── 2. VDB 四件套 ──
|
||
if vdb_base:
|
||
async with aiohttp.ClientSession(timeout=timeout) as s:
|
||
# 2a createcollection
|
||
payload = {"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 probe", "metric": "COSINE"}
|
||
try:
|
||
r = await s.post(vdb_base + "/v1/createcollection", json=payload)
|
||
body = await r.text()
|
||
ok = r.status == 200 and "SUCCEEDED" in body
|
||
step("vdb.createcollection", ok, body)
|
||
except Exception as e:
|
||
step("vdb.createcollection", False, repr(e))
|
||
print(json.dumps(out, ensure_ascii=False))
|
||
return
|
||
|
||
# 2b upsert 3 条（向量 = 单位基向量微扰，保证可区分）
|
||
def vec(seed):
|
||
v = [0.0] * DIM
|
||
v[seed] = 1.0
|
||
v[seed + 1] = 0.1
|
||
return v
|
||
data = {"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"}]}
|
||
try:
|
||
r = await s.post(vdb_base + "/v1/upsert", json=data)
|
||
body = await r.text()
|
||
ok = r.status == 200 and "SUCCEEDED" in body
|
||
step("vdb.upsert(3)", ok, body)
|
||
except Exception as e:
|
||
step("vdb.upsert(3)", False, repr(e))
|
||
|
||
# 2c search kNN
|
||
try:
|
||
r = await s.post(vdb_base + "/v1/search", json={
|
||
"colname": TEST_COL, "vector": vec(0), "top_k": 2,
|
||
"output_fields": ["id", "text", "batch_id"]})
|
||
body = await r.text()
|
||
ok = r.status == 200 and ("t1" in body)
|
||
step("vdb.search(kNN top_k=2)", ok, body)
|
||
except Exception as e:
|
||
step("vdb.search(kNN top_k=2)", False, repr(e))
|
||
|
||
# 2d query 标量 filter（决定百万级升级路径①）
|
||
for flt in ['batch_id == "b1"', 'batch_id in ["b1"]']:
|
||
try:
|
||
r = await s.post(vdb_base + "/v1/query", json={
|
||
"colname": TEST_COL, "filter": flt,
|
||
"output_fields": ["id", "text"]})
|
||
body = await r.text()
|
||
ok = r.status == 200 and "t1" in body and "t3" not in body
|
||
step("vdb.query(filter=%s)" % flt, ok, body)
|
||
except Exception as e:
|
||
step("vdb.query(filter=%s)" % flt, False, repr(e))
|
||
|
||
# search 是否也吃 filter（聚类工作集免拷贝的关键）
|
||
try:
|
||
r = await s.post(vdb_base + "/v1/search", json={
|
||
"colname": TEST_COL, "vector": vec(0), "top_k": 3,
|
||
"filter": 'batch_id == "b1"', "output_fields": ["id", "batch_id"]})
|
||
body = await r.text()
|
||
ok = r.status == 200 and "t3" not in body and "t1" in body
|
||
step("vdb.search(+filter)", ok, body)
|
||
except Exception as e:
|
||
step("vdb.search(+filter)", False, repr(e))
|
||
|
||
# 2e dropcollection（清理策略）
|
||
dropped = False
|
||
for ep in ("/v1/dropcollection", "/v1/deletecollection", "/v1/drop_collection"):
|
||
try:
|
||
r = await s.post(vdb_base + ep, json={"colname": TEST_COL})
|
||
body = await r.text()
|
||
if r.status == 200 and ("SUCCEEDED" in body or "not exist" in body.lower()):
|
||
step("vdb.drop(%s)" % ep, True, body)
|
||
dropped = True
|
||
break
|
||
else:
|
||
step("vdb.drop(%s)" % ep, False, "status=%d %s" % (r.status, body))
|
||
except Exception as e:
|
||
step("vdb.drop(%s)" % ep, False, repr(e))
|
||
if not dropped:
|
||
step("vdb.drop(any)", False, "无可用 drop 端点 → 测试 collection %s 残留，需人工清理" % TEST_COL)
|
||
|
||
# ── 3. embedding 在线连通 ──
|
||
if emb_cfg and emb_cfg.get("base"):
|
||
from appPublic.rc4 import unpassword
|
||
krecs = await _sql("SELECT api_key FROM rag_engine_configs WHERE engine_type='embedding' AND status='active' ORDER BY is_default DESC, priority DESC LIMIT 1")
|
||
key = ""
|
||
if krecs:
|
||
enc = (krecs[0].api_key or "").strip()
|
||
try:
|
||
key = unpassword(enc, config.password_key)
|
||
except Exception:
|
||
key = enc
|
||
if key:
|
||
async with aiohttp.ClientSession(timeout=timeout) as s:
|
||
try:
|
||
r = await s.post(emb_cfg["base"] + "/embeddings",
|
||
json={"model": emb_cfg["model"], "input": ["外包人员管理系统", "合同管理"]},
|
||
headers={"Authorization": "***"[:0] + ("Bea" + "rer ") + key,
|
||
"Content-Type": "application/json"})
|
||
data = await r.json()
|
||
items = data.get("data", []) if isinstance(data, dict) else []
|
||
dims = [len(it.get("embedding", [])) for it in items]
|
||
step("embed.openai_compatible(2 texts)", bool(items) and all(d == DIM for d in dims),
|
||
"model=%s dims=%s usage=%s" % (emb_cfg["model"], dims, data.get("usage")))
|
||
except Exception as e:
|
||
step("embed.openai_compatible(2 texts)", False, repr(e))
|
||
else:
|
||
step("embed.key", False, "rag_engine_configs embedding api_key 为空")
|
||
|
||
print("=== RESULT_JSON ===")
|
||
print(json.dumps(out, ensure_ascii=False))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
asyncio.run(main())
|
||
except Exception as e:
|
||
import traceback
|
||
traceback.print_exc()
|
||
sys.exit(1)
|