""" RAG 文件处理管线 — 文档解析 / 人脸检测 / 声纹处理 GPU 服务映射: 人脸: https://media.opencomputing.net:10443/face/api/detect 声纹: https://media.opencomputing.net:10443/voiceprint/extract/submit """ import json, os, uuid, base64 # File type classification TEXT_EXTS = {'.txt', '.md', '.csv', '.json', '.xml', '.html', '.htm', '.py', '.js', '.css', '.yaml', '.yml', '.log', '.rst'} DOC_EXTS = {'.pdf', '.docx', '.pptx', '.xlsx', '.doc', '.ppt', '.xls'} IMAGE_EXTS = {'.jpg', '.jpeg', '.png', '.bmp', '.gif', '.webp'} AUDIO_EXTS = {'.mp3', '.wav', '.flac', '.ogg', '.m4a', '.aac'} VIDEO_EXTS = {'.mp4', '.avi', '.mov', '.mkv', '.webm'} # GPU service endpoints FACE_DETECT_URL = "https://media.opencomputing.net:10443/face/api/detect" FACE_RECOGNIZE_URL = "https://media.opencomputing.net:10443/face/api/recognize" VOICEPRINT_EXTRACT_URL = "https://media.opencomputing.net:10443/voiceprint/extract/submit" VOICEPRINT_STATUS_URL = "https://media.opencomputing.net:10443/voiceprint/extract/status" def classify_file(file_name): ext = os.path.splitext(file_name)[1].lower() if ext in TEXT_EXTS: return 'text' if ext in DOC_EXTS: return 'document' if ext in IMAGE_EXTS: return 'image' if ext in AUDIO_EXTS: return 'audio' if ext in VIDEO_EXTS: return 'video' return 'other' async def extract_text(file_data, file_name): ext = os.path.splitext(file_name)[1].lower() text = "" if ext in TEXT_EXTS: text = file_data.decode('utf-8', errors='replace') elif ext == '.pdf': try: import io from PyPDF2 import PdfReader reader = PdfReader(io.BytesIO(file_data)) text = '\n'.join(p.extract_text() or '' for p in reader.pages) except Exception as e: text = "[PDF error: " + str(e) + "]" elif ext == '.docx': try: import io from docx import Document doc = Document(io.BytesIO(file_data)) text = '\n'.join(p.text for p in doc.paragraphs if p.text.strip()) except Exception as e: text = "[DOCX error: " + str(e) + "]" elif ext == '.pptx': try: import io from pptx import Presentation prs = Presentation(io.BytesIO(file_data)) slides = [] for slide in prs.slides: parts = [s.text for s in slide.shapes if hasattr(s, 'text') and s.text.strip()] slides.append(' '.join(parts)) text = '\n'.join(slides) except Exception as e: text = "[PPTX error: " + str(e) + "]" elif ext in {'.xlsx', '.xls'}: try: import io from openpyxl import load_workbook wb = load_workbook(io.BytesIO(file_data), read_only=True, data_only=True) rows = [] for ws in wb.worksheets: for row in ws.iter_rows(values_only=True): r = ' | '.join(str(c) if c is not None else '' for c in row) if r.strip(): rows.append(r) text = '\n'.join(rows) except Exception as e: text = "[XLSX error: " + str(e) + "]" return text.strip() async def ingest_to_rag(env, text, kb_id, doc_id): try: from pipeline import ingest as pipeline_ingest result = pipeline_ingest( text, pipeline_name="kg-rag-standard", collection=kb_id, graph_name=kb_id, llm_func=None) return result except Exception as e: return {"error": str(e), "chunks": 0} async def detect_faces(file_data, file_name): """Detect faces via GPU face-service (InsightFace buffalo_l)""" import aiohttp try: b64 = base64.b64encode(file_data).decode() payload = {"image": b64} async with aiohttp.ClientSession() as session: async with session.post( FACE_DETECT_URL, json=payload, timeout=aiohttp.ClientTimeout(total=30) ) as resp: if resp.status == 200: result = await resp.json() return { "faces": result.get("faces", result.get("count", 0)), "details": result } return {"error": "face service returned " + str(resp.status), "faces": 0} except Exception as e: return {"error": str(e), "faces": 0} async def extract_voiceprint(file_data, file_name): import aiohttp try: async with aiohttp.ClientSession() as session: form = aiohttp.FormData() form.add_field('file', file_data, filename=file_name, content_type='audio/wav') async with session.post(VOICEPRINT_EXTRACT_URL, data=form, timeout=aiohttp.ClientTimeout(total=60)) as resp: if resp.status == 200: r = await resp.json() if r.get('status') == 'SUCCEEDED': return {'voiceprint': r, 'speakers': 1} return {'error': r.get('error', ''), 'speakers': 0} return {'error': 'service returned ' + str(resp.status), 'speakers': 0} except Exception as e: return {'error': str(e), 'speakers': 0} async def process_upload(env, file_data, kb_id, folder_id, file_name): """Full upload pipeline — save + classify + process + DB record""" from appPublic.uniqueID import getID doc_id = getID() ext = '.' + file_name.rsplit('.', 1)[1] if '.' in file_name else '.bin' saved_name = doc_id + ext file_path = '/d/rag/ragserver/pkgs/rag/rag/files/' + saved_name with open(file_path, 'wb') as f: f.write(file_data) file_size = len(file_data) ft = classify_file(file_name) ingest_result = None face_result = None voice_result = None status = 'pending' if ft in ('text', 'document'): text = await extract_text(file_data, file_name) if text and len(text) > 10: ingest_result = await ingest_to_rag(env, text, kb_id, doc_id) status = 'done' if ingest_result and not ingest_result.get('error') else 'error' else: status = 'done' elif ft == 'image': face_result = await detect_faces(file_data, file_name) status = 'done' elif ft == 'audio': voice_result = await extract_voiceprint(file_data, file_name) status = 'done' elif ft == 'video': face_result = await detect_faces(file_data, file_name) voice_result = await extract_voiceprint(file_data, file_name) status = 'done' else: status = 'done' import json as _json meta = _json.dumps({'file_type': ft, 'face': face_result, 'voice': voice_result, 'ingest': str(ingest_result)[:500]}, ensure_ascii=False, default=str) userorgid = await env.get_userorgid() from sqlor.dbpools import get_sor_context async with get_sor_context(env, 'rag') as sor: await sor.sqlExe( "INSERT INTO rag_documents (id, kb_id, folder_id, file_name, file_type, file_size, file_path, mime_type, status, metadata, org_id, created_at, updated_at) " "VALUES (" + "${id}$, ${kb_id}$, ${folder_id}$, ${file_name}$, 'other', ${file_size}$, ${file_path}$, 'application/octet-stream', ${status}$, ${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": "/idfile/files/" + saved_name, "status": status, "meta": meta, "org_id": userorgid}) await sor.sqlExe( "UPDATE rag_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 ingest_result and status == 'done': chunks_n = ingest_result.get('chunks', 0) if isinstance(ingest_result, dict) else 0 if chunks_n: await sor.sqlExe("UPDATE rag_documents SET chunk_count=${chunks}$ WHERE id=${id}$", {"chunks": chunks_n, "id": doc_id}) await sor.sqlExe("UPDATE rag_knowledge_bases SET chunk_count=chunk_count+${n}$ WHERE id=${kb_id}$", {"n": chunks_n, "kb_id": kb_id}) faces_n = face_result.get('faces', 0) if face_result and isinstance(face_result, dict) else 0 speakers_n = voice_result.get('speakers', 0) if voice_result and isinstance(voice_result, dict) else 0 return _json.dumps({ "status": "SUCCEEDED", "doc_id": doc_id, "file_name": file_name, "file_size": file_size, "folder_id": folder_id, "file_type": ft, "rag_status": status, "faces": faces_n, "speakers": speakers_n }, ensure_ascii=False, default=str)