add text extraction for PDF/DOCX/PPTX/XLSX and MP4 keyframe CLIP embedding in upload_file.dspy

This commit is contained in:
yumoqing 2026-07-30 18:06:26 +08:00
parent c6110522c1
commit 1cb82b912d

View File

@ -1,3 +1,4 @@
import base64, os, subprocess
ns = params_kw.copy()
kb_id = ns.get('kb_id', '')
folder_id = ns.get('folder', '')
@ -32,6 +33,33 @@ 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')
@ -83,6 +111,30 @@ if ext_l in video_exts:
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