yumoqing/rag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(embedding): 知识库支持文本(bge-m3 /txte)/多媒体(CLIP /mme)双向量引擎

Browse Source 
- create_kb/new_kb_form 加向量引擎选择(文本 bge-m3 / 多媒体 clip-vith14)
- upload_file/batch_ingest/search_result/init.py 按 embedding_engine 路由到 /txte 或 /mme
- 文本知识库上传媒体文件友好拒绝(不再崩溃报错)
This commit is contained in:
ymq 2026-08-25 13:15:11 +08:00
parent 5354108546
commit 1054dbb3a9
7 changed files with 107 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

49
rag/init.py
View File

@ -113,7 +113,7 @@ async def search_handler(request, params_kw, *args, **kwargs):
# Process query + media → embedding vector
query_vec = None
if query or file_data:
query_vec = await _build_search_vector(query, file_data, file_name)
query_vec = await _build_search_vector(query, file_data, file_name, env, kb_id)
if not query_vec:
return json.dumps({"status": "SUCCEEDED", "data": {"results": [], "total": 0, "message": "no query or file provided"}}, ensure_ascii=False)
@ -181,8 +181,8 @@ async def _resolve_search_kbs(env, userorgid, kb_id):
return [r.id for r in recs]
async def _build_search_vector(query, file_data, file_name):
"""Build search embedding from text query + media file"""
async def _build_search_vector(query, file_data, file_name, env=None, kb_id=''):
"""Build search embedding from text query + media file（按 kb 向量引擎选端点）"""
texts = []
if query:
texts.append(query)
@ -216,11 +216,23 @@ async def _build_search_vector(query, file_data, file_name):
combined = " ".join(texts)
try:
resp = await _call_uapi("rag-embedding", "embed", {
"texts": [combined],
"model": "CLIP-ViT-H-14"
})
vecs = resp.get("embeddings", []) if isinstance(resp, dict) else []
emb_engine = 'clip-vith14'
if env is not None and kb_id:
async with get_sor_context(env, 'rag') as sor:
krecs = await sor.sqlExe("SELECT embedding_engine FROM knowledge_bases WHERE id=${kb_id}$", {"kb_id": kb_id})
if krecs:
emb_engine = (getattr(krecs[0], 'embedding_engine', '') or 'clip-vith14').strip()
if emb_engine == 'bge-m3':
emb_url = 'https://embedding.opencomputing.net:10443/txte/api/embed'
emb_model = 'bge-m3'
else:
emb_url = 'https://embedding.opencomputing.net:10443/mme/api/embed'
emb_model = 'CLIP-ViT-H-14'
import aiohttp
async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=15)) as s:
async with s.post(emb_url, json={"texts": [combined], "model": emb_model}) as resp:
emb_resp = await resp.json()
vecs = emb_resp.get("text_embeddings", emb_resp.get("embeddings", [])) if isinstance(emb_resp, dict) else []
return vecs[0] if vecs else None
except Exception as e:
exception(f"query embedding failed: {e}")
@ -480,11 +492,24 @@ async def _rag_ingest_async(env, text, kb_id, doc_id):
return {"chunks": 0}
chunk_count = len(chunks)
# 1. Embedding
# 1. Embedding（按知识库向量引擎选文本/多模态端点）
try:
emb_resp = await _call_uapi("rag-embedding", "embed",
{"texts": chunks, "model": "CLIP-ViT-H-14"})
embeddings = emb_resp.get("embeddings", []) if isinstance(emb_resp, dict) else []
emb_engine = 'clip-vith14'
async with get_sor_context(env, 'rag') as sor:
krecs = await sor.sqlExe("SELECT embedding_engine FROM knowledge_bases WHERE id=${kb_id}$", {"kb_id": kb_id})
if krecs:
emb_engine = (getattr(krecs[0], 'embedding_engine', '') or 'clip-vith14').strip()
if emb_engine == 'bge-m3':
emb_url = 'https://embedding.opencomputing.net:10443/txte/api/embed'
emb_model = 'bge-m3'
else:
emb_url = 'https://embedding.opencomputing.net:10443/mme/api/embed'
emb_model = 'CLIP-ViT-H-14'
import aiohttp
async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=15)) as s:
async with s.post(emb_url, json={"texts": chunks, "model": emb_model}) as resp:
emb_resp = await resp.json()
embeddings = emb_resp.get("text_embeddings", emb_resp.get("embeddings", [])) if isinstance(emb_resp, dict) else []
except Exception as e:
exception(f"embedding failed: {e}")
embeddings = []

12
wwwroot/knowledge_bases_list/batch_ingest.dspy
View File

@ -9,6 +9,14 @@ async with get_sor_context(env, 'rag') as sor:
"SELECT id, file_name, file_size, file_path FROM documents WHERE kb_id=${kb_id}$ AND (chunk_count IS NULL OR chunk_count=0 OR status='pending') LIMIT 20",
{"kb_id": kb_id})
rows = [dict(r) for r in recs]
krecs = await sor.sqlExe("SELECT embedding_engine FROM knowledge_bases WHERE id=${kb_id}$", {"kb_id": kb_id})
emb_engine = (getattr(krecs[0], 'embedding_engine', '') or 'clip-vith14').strip() if krecs else 'clip-vith14'
if emb_engine == 'bge-m3':
emb_url = 'https://embedding.opencomputing.net:10443/txte/api/embed'
emb_model = 'bge-m3'
else:
emb_url = 'https://embedding.opencomputing.net:10443/mme/api/embed'
emb_model = 'CLIP-ViT-H-14'
if not rows:
return json.dumps({"widgettype": "Text", "options": {"text": "No unprocessed documents found"}}, ensure_ascii=False)
@ -173,8 +181,8 @@ async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=60)) as ses
# Embed
try:
r2 = await session.post('https://embedding.opencomputing.net:10443/api/embed',
json={"texts": chunks, "model": "CLIP-ViT-H-14"})
r2 = await session.post(emb_url,
json={"texts": chunks, "model": emb_model})
emb = await r2.json() if r2.status == 200 else {}
embeddings = emb.get("text_embeddings", emb.get("embeddings", []))
except:

8
wwwroot/knowledge_bases_list/create_kb.dspy
View File

@ -7,10 +7,14 @@ async with db.sqlorContext(dbname) as sor:
kb_id = uuid()
name = ns.get("name", "")
desc = ns.get("description", "")
# 向量引擎：clip-vith14=多媒体(CLIP)，bge-m3=文本(bge-m3)。默认 clip-vith14 兼容旧数据
emb_type = ns.get("embedding_type", "") or ns.get("embedding_engine", "")
if emb_type not in ("bge-m3", "clip-vith14"):
emb_type = "clip-vith14"
await sor.sqlExe(
"INSERT INTO knowledge_bases (id, name, description, org_id, embedding_engine, vdb_collection, doc_count, total_size, chunk_count, status, created_at) "
"VALUES (${id}$, ${name}$, ${desc}$, ${org_id}$, 'CLIP ViT-H-14', 'rag_collection', 0, 0, 0, 'active', NOW())",
{"id": kb_id, "name": name, "desc": desc, "org_id": userorgid})
"VALUES (${id}$, ${name}$, ${desc}$, ${org_id}$, ${emb}$, 'rag_collection', 0, 0, 0, 'active', NOW())",
{"id": kb_id, "name": name, "desc": desc, "org_id": userorgid, "emb": emb_type})
return {
"widgettype": "urlwidget",
"options": {"url": entire_url('/rag/knowledge_bases_list/index.ui')}

5
wwwroot/knowledge_bases_list/embedding_options.dspy Normal file
View File

@ -0,0 +1,5 @@
import json
return json.dumps([
{"value": "clip-vith14", "text": "多媒体（图文音视）"},
{"value": "bge-m3", "text": "文本（文档检索）"}
], ensure_ascii=False)

3
wwwroot/knowledge_bases_list/new_kb_form.ui
View File

@ -10,7 +10,8 @@
"cols": 1,
"fields": [
{"name": "name", "label": "名称", "uitype": "str", "required": true},
{"name": "description", "label": "描述", "uitype": "text"}
{"name": "description", "label": "描述", "uitype": "text"},
{"name": "embedding_type", "label": "向量引擎", "uitype": "code", "value": "clip-vith14", "dataurl": "{{entire_url('./embedding_options.dspy')}}"}
]
}
}

19
wwwroot/knowledge_bases_list/search_result.dspy
View File

@ -75,13 +75,28 @@ hits = []
raw_rows = []
kw_rows = []
# 读知识库向量引擎，决定用文本(bge-m3 /txte)还是多媒体(CLIP /mme) embedding
emb_engine = 'clip-vith14'
try:
async with get_sor_context(env, 'rag') as sor:
krecs = await sor.sqlExe("SELECT embedding_engine FROM knowledge_bases WHERE id=${kb_id}$", {"kb_id": kb_id})
if krecs:
emb_engine = (getattr(krecs[0], 'embedding_engine', '') or 'clip-vith14').strip()
except: pass
if emb_engine == 'bge-m3':
emb_url = 'https://embedding.opencomputing.net:10443/txte/api/embed'
emb_model = 'bge-m3'
else:
emb_url = 'https://embedding.opencomputing.net:10443/mme/api/embed'
emb_model = 'CLIP-ViT-H-14'
if query:
# 1. Embed
vec = []
try:
client = StreamHttpClient()
resp = await client.request('POST', 'https://embedding.opencomputing.net:10443/api/embed',
json={"texts": [query], "model": "CLIP-ViT-H-14"})
resp = await client.request('POST', emb_url,
json={"texts": [query], "model": emb_model})
emb = json.loads(resp)
vec = emb.get("text_embeddings", emb.get("embeddings", [[]]))[0]
except:

34
wwwroot/knowledge_bases_list/upload_file.dspy
View File

@ -48,6 +48,22 @@ ext_l = ext.lower()
async def ingest_doc(doc_id, kb_id, file_name, ext_l, real_path):
db = DBPools()
# 读知识库向量引擎：bge-m3=文本(走 /txte)，clip-vith14=多媒体(走 /mme)
emb_engine = 'clip-vith14'
try:
async with db.sqlorContext('rag') as sor:
krecs = await sor.sqlExe("SELECT embedding_engine FROM knowledge_bases WHERE id=${kb_id}$", {"kb_id": kb_id})
if krecs:
emb_engine = (getattr(krecs[0], 'embedding_engine', '') or 'clip-vith14').strip()
except: pass
is_text = emb_engine == 'bge-m3'
if is_text:
emb_url = 'https://embedding.opencomputing.net:10443/txte/api/embed'
emb_model = 'bge-m3'
else:
emb_url = 'https://embedding.opencomputing.net:10443/mme/api/embed'
emb_model = 'CLIP-ViT-H-14'
async def ensure_vdb_collection(client, kb_id):
payload = {"colname": kb_id, "fields": [{"name": "id", "type": "str", "is_primary": True, "max_length": 64}, {"name": "vector", "type": "fvector", "dim": 1024}, {"name": "text", "type": "str", "max_length": 65535}], "description": "RAG kb", "metric": "COSINE"}
await client.request('POST', 'https://vectordb.opencomputing.net:10443/v1/createcollection', json=payload)
@ -97,6 +113,16 @@ async def ingest_doc(doc_id, kb_id, file_name, ext_l, real_path):
if ext_l in text_exts:
text = file_data.decode('utf-8', errors='replace')
# --- 文本知识库不支持媒体文件 ---
if is_text and (ext_l in image_exts or ext_l in audio_exts or ext_l in video_exts):
try:
async with db.sqlorContext('rag') as sor:
await sor.sqlExe(
"UPDATE documents SET status='failed', metadata=${meta}$, updated_at=NOW() WHERE id=${id}$",
{"id": doc_id, "meta": json.dumps({"error": "文本知识库不支持媒体文件，请上传文本类文件(txt/md/pdf/docx等)或改用多媒体知识库"}, ensure_ascii=False)})
except: pass
return
# --- IMAGE: face detection ---
if ext_l in image_exts:
img_b64 = base64.b64encode(file_data).decode()
@ -150,8 +176,8 @@ async def ingest_doc(doc_id, kb_id, file_name, ext_l, real_path):
# --- CLIP image embedding for video frame ---
try:
client2 = StreamHttpClient()
resp2 = await client2.request('POST', 'https://embedding.opencomputing.net:10443/api/embed',
json={"images": [img_b64], "model": "CLIP-ViT-H-14"})
resp2 = await client2.request('POST', emb_url,
json={"images": [img_b64], "model": emb_model})
emb_data = json.loads(resp2)
img_embeddings = emb_data.get("image_embeddings", emb_data.get("embeddings", []))
except:
@ -226,8 +252,8 @@ async def ingest_doc(doc_id, kb_id, file_name, ext_l, real_path):
if chunks:
try:
client = StreamHttpClient()
resp = await client.request('POST', 'https://embedding.opencomputing.net:10443/api/embed',
json={"texts": chunks, "model": "CLIP-ViT-H-14"})
resp = await client.request('POST', emb_url,
json={"texts": chunks, "model": emb_model})
emb_data = json.loads(resp)
embeddings = emb_data.get("text_embeddings", emb_data.get("embeddings", []))
except: