212 lines
9.7 KiB
Plaintext
212 lines
9.7 KiB
Plaintext
import aiohttp, base64, os, subprocess
|
|
ns = params_kw.copy()
|
|
kb_id = ns.get('kb_id', '')
|
|
folder_id = ns.get('folder', '')
|
|
file_name = ns.get('file_name', 'upload.bin')
|
|
if not kb_id:
|
|
return json.dumps({"status": "error", "error": "kb_id required"}, ensure_ascii=False)
|
|
|
|
file_data = await request.read()
|
|
if not file_data:
|
|
return json.dumps({"status": "error", "error": "no file data"}, ensure_ascii=False)
|
|
|
|
env = request._run_ns
|
|
userorgid = await env.get_userorgid()
|
|
|
|
# Save file via FileStorage (returns web path e.g. /idfile/191/193/197/97/xxx.txt)
|
|
web_path = await env.save_file(file_data, file_name)
|
|
real_path = env.realpath(web_path)
|
|
|
|
doc_id = str(uuid()).replace('-', '')[:16]
|
|
file_size = len(file_data)
|
|
ext = '.' + file_name.rsplit('.', 1)[1] if '.' in file_name else '.bin'
|
|
ext_l = ext.lower()
|
|
|
|
text_exts = {'.txt', '.md', '.csv', '.json', '.xml', '.html', '.htm', '.py', '.js', '.css', '.yaml', '.yml', '.log', '.rst'}
|
|
image_exts = {'.jpg', '.jpeg', '.png', '.bmp', '.gif', '.webp'}
|
|
audio_exts = {'.mp3', '.wav', '.flac', '.ogg', '.m4a', '.aac'}
|
|
video_exts = {'.mp4', '.avi', '.mov', '.mkv', '.webm'}
|
|
|
|
text = ''
|
|
chunks_n = 0
|
|
face_count = 0
|
|
voice_speakers = 0
|
|
meta_parts = {}
|
|
|
|
# --- OFFICE DOCS: text extraction (PDF/DOCX/PPTX/XLSX) ---
|
|
if ext_l == '.pdf' and not text:
|
|
import io; from PyPDF2 import PdfReader
|
|
reader = PdfReader(io.BytesIO(file_data))
|
|
text = '\n'.join(p.extract_text() or '' for p in reader.pages)
|
|
elif ext_l == '.docx' and not text:
|
|
import io; from docx import Document
|
|
doc = Document(io.BytesIO(file_data))
|
|
text = '\n'.join(p.text for p in doc.paragraphs)
|
|
elif ext_l == '.pptx' and not text:
|
|
import io; from pptx import Presentation
|
|
prs = Presentation(io.BytesIO(file_data))
|
|
parts = []
|
|
for slide in prs.slides:
|
|
for shape in slide.shapes:
|
|
if hasattr(shape, 'text') and shape.text:
|
|
parts.append(shape.text)
|
|
text = '\n'.join(parts)
|
|
elif ext_l == '.xlsx' and not text:
|
|
import io; from openpyxl import load_workbook
|
|
wb = load_workbook(io.BytesIO(file_data), data_only=True)
|
|
parts = []
|
|
for sheet in wb.worksheets:
|
|
for row in sheet.iter_rows(values_only=True):
|
|
parts.append('\t'.join(str(c or '') for c in row))
|
|
text = '\n'.join(parts)
|
|
|
|
# --- TEXT EXTRACTION ---
|
|
if ext_l in text_exts:
|
|
text = file_data.decode('utf-8', errors='replace')
|
|
|
|
# --- IMAGE: face detection ---
|
|
if ext_l in image_exts:
|
|
img_b64 = base64.b64encode(file_data).decode()
|
|
async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=15)) as s:
|
|
try:
|
|
r = await s.post('https://media.opencomputing.net/face/api/detect', json={"images": [img_b64]})
|
|
if r.status == 200:
|
|
fd = await r.json()
|
|
results = fd.get("results", [])
|
|
if results and isinstance(results[0], dict):
|
|
face_count = len(results[0].get("faces", results[0].get("detections", [])))
|
|
meta_parts['face'] = face_count
|
|
except: pass
|
|
|
|
# --- AUDIO: voiceprint ---
|
|
if ext_l in audio_exts:
|
|
async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=60)) as s:
|
|
try:
|
|
form = aiohttp.FormData()
|
|
form.add_field('file', file_data, filename=file_name)
|
|
r = await s.post('https://media.opencomputing.net/voiceprint/extract/submit', data=form)
|
|
if r.status == 200:
|
|
vd = await r.json()
|
|
voice_speakers = vd.get('speakers', 1) if vd.get('status') == 'SUCCEEDED' else (1 if vd.get('embedding') else 0)
|
|
meta_parts['voiceprint'] = voice_speakers
|
|
except: pass
|
|
|
|
# --- VIDEO: frame extraction ---
|
|
if ext_l in video_exts:
|
|
meta_parts['video'] = 'pending'
|
|
try:
|
|
tmp_img = '/tmp/' + doc_id + '_frame.jpg'
|
|
subprocess.run(['ffmpeg', '-y', '-i', real_path, '-vframes', '1', '-q:v', '2', tmp_img],
|
|
capture_output=True, timeout=30)
|
|
if os.path.exists(tmp_img):
|
|
with open(tmp_img, 'rb') as fi:
|
|
frame_data = fi.read()
|
|
img_b64 = base64.b64encode(frame_data).decode()
|
|
async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=15)) as s:
|
|
try:
|
|
r = await s.post('https://media.opencomputing.net/face/api/detect', json={"images": [img_b64]})
|
|
if r.status == 200:
|
|
fd = await r.json()
|
|
results = fd.get("results", [])
|
|
if results and isinstance(results[0], dict):
|
|
face_count = len(results[0].get("faces", results[0].get("detections", [])))
|
|
except: pass
|
|
# --- CLIP image embedding for video frame ---
|
|
if os.path.exists(tmp_img):
|
|
img_b64_frame = base64.b64encode(frame_data).decode()
|
|
try:
|
|
r2 = await s.post('https://embedding.opencomputing.net/api/embed',
|
|
json={"images": [img_b64_frame], "model": "CLIP-ViT-H-14"})
|
|
emb_data = await r2.json() if r2.status == 200 else {}
|
|
img_embeddings = emb_data.get("image_embeddings", emb_data.get("embeddings", []))
|
|
except:
|
|
img_embeddings = []
|
|
if img_embeddings:
|
|
try:
|
|
vdb_data = {"colname": kb_id, "data": [
|
|
{"id": doc_id + "_frame0", "vector": img_embeddings[0], "text": file_name}
|
|
]}
|
|
await s.post('https://vectordb.opencomputing.net/v1/upsert', json=vdb_data)
|
|
async with get_sor_context(env, 'rag') as sor:
|
|
await sor.sqlExe(
|
|
"INSERT INTO document_chunks (id, doc_id, kb_id, chunk_index, content, vector_id, created_at) "
|
|
"VALUES (${id}$, ${doc_id}$, ${kb_id}$, 0, ${content}$, ${vid}$, NOW())",
|
|
{"id": doc_id + "_c0", "doc_id": doc_id, "kb_id": kb_id,
|
|
"content": file_name, "vid": doc_id + "_frame0"})
|
|
except:
|
|
pass
|
|
os.remove(tmp_img)
|
|
except: pass
|
|
|
|
# --- RAG INGEST for text ---
|
|
if text and len(text.strip()) > 10:
|
|
paragraphs = text.split('\n')
|
|
chunks = []
|
|
cur = ''
|
|
for p in paragraphs:
|
|
p = p.strip()
|
|
if not p:
|
|
if cur: chunks.append(cur); cur = ''
|
|
continue
|
|
if len(cur) + len(p) < 500:
|
|
cur = (cur + '\n' + p).strip()
|
|
else:
|
|
if cur: chunks.append(cur)
|
|
cur = p
|
|
if cur: chunks.append(cur)
|
|
|
|
if chunks:
|
|
async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=15)) as s:
|
|
try:
|
|
r = await s.post('https://embedding.opencomputing.net/api/embed',
|
|
json={"texts": chunks, "model": "CLIP-ViT-H-14"})
|
|
emb_data = await r.json() if r.status == 200 else {}
|
|
embeddings = emb_data.get("text_embeddings", emb_data.get("embeddings", []))
|
|
except:
|
|
embeddings = []
|
|
|
|
vector_ids = []
|
|
if embeddings:
|
|
try:
|
|
vdb_data = {"colname": kb_id, "data": [
|
|
{"id": doc_id + "_" + str(i), "vector": emb, "text": chunks[i]}
|
|
for i, emb in enumerate(embeddings)]}
|
|
await s.post('https://vectordb.opencomputing.net/v1/upsert', json=vdb_data)
|
|
vector_ids = [doc_id + "_" + str(i) for i in range(len(embeddings))]
|
|
except:
|
|
pass
|
|
|
|
async with get_sor_context(env, 'rag') as sor:
|
|
for i, chunk_text in enumerate(chunks):
|
|
vid = vector_ids[i] if i < len(vector_ids) else ''
|
|
await sor.sqlExe(
|
|
"INSERT INTO document_chunks (id, doc_id, kb_id, chunk_index, content, vector_id, created_at) "
|
|
"VALUES (${id}$, ${doc_id}$, ${kb_id}$, ${idx}$, ${content}$, ${vid}$, NOW())",
|
|
{"id": doc_id + "_c" + str(i), "doc_id": doc_id, "kb_id": kb_id,
|
|
"idx": i, "content": chunk_text[:2000], "vid": vid})
|
|
chunks_n = len(chunks)
|
|
|
|
# --- SAVE TO DB ---
|
|
meta_json = json.dumps(meta_parts, ensure_ascii=False)
|
|
status = 'done'
|
|
async with get_sor_context(env, 'rag') as sor:
|
|
await sor.sqlExe(
|
|
"INSERT INTO documents (id, kb_id, folder_id, file_name, file_type, file_size, file_path, mime_type, status, chunk_count, metadata, org_id, created_at, updated_at) "
|
|
"VALUES (${id}$, ${kb_id}$, ${folder_id}$, ${file_name}$, 'other', ${file_size}$, ${file_path}$, 'application/octet-stream', ${status}$, ${chunks}$, ${meta}$, ${org_id}$, NOW(), NOW())",
|
|
{"id": doc_id, "kb_id": kb_id, "folder_id": folder_id, "file_name": file_name,
|
|
"file_size": file_size, "file_path": web_path,
|
|
"status": status, "chunks": chunks_n, "meta": meta_json, "org_id": userorgid})
|
|
await sor.sqlExe(
|
|
"UPDATE knowledge_bases SET doc_count=doc_count+1, total_size=total_size+${size}$ WHERE id=${kb_id}$",
|
|
{"size": file_size, "kb_id": kb_id})
|
|
if chunks_n:
|
|
await sor.sqlExe(
|
|
"UPDATE knowledge_bases SET chunk_count=chunk_count+${n}$ WHERE id=${kb_id}$",
|
|
{"n": chunks_n, "kb_id": kb_id})
|
|
|
|
result = {"status": "SUCCEEDED", "doc_id": doc_id, "file_name": file_name,
|
|
"file_size": file_size, "folder_id": folder_id,
|
|
"text_len": len(text), "chunks": chunks_n,
|
|
"faces": face_count, "speakers": voice_speakers}
|
|
return json.dumps(result, ensure_ascii=False, default=str)
|