From c25351a65e027407a6320d962ea8cfdbf44e4187 Mon Sep 17 00:00:00 2001 From: yumoqing Date: Wed, 22 Jul 2026 18:09:58 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20file=20upload=20=E2=86=92=20RAG=20inges?= =?UTF-8?q?t=20+=20file=20delete=20=E2=86=92=20cleanup=20pipeline?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- conf/config.json | 3 +- init.py | 217 ++++++++++++++++++++++++++++++++++++++++------- 2 files changed, 189 insertions(+), 31 deletions(-) diff --git a/conf/config.json b/conf/config.json index 41ab30a..7222612 100644 --- a/conf/config.json +++ b/conf/config.json @@ -51,7 +51,8 @@ {"leading": "/api/kb/list", "registerfunction": "kb_list"}, {"leading": "/api/engines", "registerfunction": "engines"}, {"leading": "/api/search", "registerfunction": "search"}, - {"leading": "/api/ingest", "registerfunction": "ingest"} + {"leading": "/api/doc/upload", "registerfunction": "doc_upload"}, + {"leading": "/api/doc/delete", "registerfunction": "doc_delete"} ] }, "hot_reload": true diff --git a/init.py b/init.py index 4aa987d..e271704 100644 --- a/init.py +++ b/init.py @@ -1,44 +1,37 @@ # -*- coding:utf-8 -*- """ -RagServer DSPY Handlers - RAG 核心业务逻辑 +RagServer DSPY Handlers - RAG 核心业务逻辑 + 文件上传/删除 """ from traceback import format_exc from ahserver.serverenv import ServerEnv from appPublic.registerfunction import RegisterFunction from appPublic.log import debug, exception from sqlor.dbpools import get_sor_context -import json +import json, os, uuid, time async def status_handler(request, params_kw, *args, **kwargs): - """服务状态""" return json.dumps({ - "service": "ragserver", - "version": "0.1.0", - "endpoints": [ - "/api/status", "/api/kb/list", "/api/doc/upload", - "/api/doc/ingest", "/api/search", "/api/engines" - ] + "service": "ragserver", "version": "0.1.0", + "endpoints": ["/api/status", "/api/kb/list", "/api/doc/upload", + "/api/doc/delete", "/api/search", "/api/engines"] }, indent=2, ensure_ascii=False) async def kb_list_handler(request, params_kw, *args, **kwargs): - """列出当前用户的知识库""" env = request._run_ns try: userorgid = await env.get_userorgid() async with get_sor_context(env, 'rag') as sor: recs = await sor.R("knowledge_bases", {"org_id": userorgid}) rows = [dict(r) for r in recs] - result = {"status": "SUCCEEDED", "data": {"rows": rows, "total": len(rows)}} - return json.dumps(result, ensure_ascii=False, default=str) + return json.dumps({"status": "SUCCEEDED", "data": {"rows": rows, "total": len(rows)}}, ensure_ascii=False, default=str) except Exception as e: exception(f"kb_list: {e}, {format_exc()}") return json.dumps({"error": str(e)}) async def engines_handler(request, params_kw, *args, **kwargs): - """列出可用引擎配置""" env = request._run_ns try: userorgid = await env.get_userorgid() @@ -53,7 +46,6 @@ async def engines_handler(request, params_kw, *args, **kwargs): async def search_handler(request, params_kw, *args, **kwargs): - """混合检索""" env = request._run_ns try: query = params_kw.get("query", "") @@ -75,33 +67,198 @@ async def search_handler(request, params_kw, *args, **kwargs): return json.dumps({"error": str(e)}) -async def ingest_handler(request, params_kw, *args, **kwargs): - """文档入库""" +async def doc_upload_handler(request, params_kw, *args, **kwargs): + """文件上传 → 保存 → DB记录 → 触发RAG入库""" env = request._run_ns try: - document = params_kw.get("document", "") + userorgid = await env.get_userorgid() kb_id = params_kw.get("kb_id", "") - if not document: - return json.dumps({"error": "document text required"}) - try: - from pipeline import ingest as pipeline_ingest - result = pipeline_ingest(document, pipeline_name="kg-rag-standard", - collection=kb_id or "knowledge", - graph_name=kb_id or "knowledge", llm_func=None) - return json.dumps(result, ensure_ascii=False) - except ImportError: - return json.dumps({"status": "FALLBACK", "message": "pipeline not available"}) + folder = params_kw.get("folder", "/") + if not kb_id: + return json.dumps({"error": "kb_id required"}) + + # Read uploaded file from request + reader = await request.multipart() + file_data = None + file_name = None + content_type = None + while True: + part = await reader.next() + if part is None: + break + if part.name == "file": + file_name = part.filename or "upload.bin" + file_data = await part.read() + content_type = part.headers.get("Content-Type", "application/octet-stream") + + if not file_data: + return json.dumps({"error": "no file uploaded"}) + + # Save file + doc_id = uuid.uuid4().hex[:16] + ext = os.path.splitext(file_name)[1] or ".bin" + saved_name = f"{doc_id}{ext}" + files_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "files") + os.makedirs(files_dir, exist_ok=True) + file_path = os.path.join(files_dir, saved_name) + with open(file_path, "wb") as f: + f.write(file_data) + + file_size = len(file_data) + file_type = _detect_file_type(file_name, content_type) + + # Create document record + async with get_sor_context(env, 'rag') as sor: + await sor.sqlExe( + "INSERT INTO documents (id, kb_id, file_name, file_type, file_size, file_path, mime_type, status, org_id, created_at, updated_at) " + "VALUES (${id}$, ${kb_id}$, ${file_name}$, ${file_type}$, ${file_size}$, ${file_path}$, ${mime_type}$, 'pending', ${org_id}$, NOW(), NOW())", + {"id": doc_id, "kb_id": kb_id, "file_name": file_name, "file_type": file_type, + "file_size": file_size, "file_path": "/idfile/files/" + saved_name, + "mime_type": content_type, "org_id": userorgid}) + + # Update KB stats + 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}) + + # Trigger async ingest for text-based files + ingest_result = None + if file_type == "text": + try: + text = file_data.decode("utf-8", errors="replace") + from pipeline import ingest as pipeline_ingest + ingest_result = pipeline_ingest( + text, pipeline_name="kg-rag-standard", + collection=kb_id, graph_name=kb_id, llm_func=None) + # Update status to done + async with get_sor_context(env, 'rag') as sor: + await sor.sqlExe( + "UPDATE documents SET status='done', chunk_count=${chunks}$ WHERE id=${id}$", + {"chunks": ingest_result.get("chunks", 0) if ingest_result else 0, "id": doc_id}) + if ingest_result: + await sor.sqlExe( + "UPDATE knowledge_bases SET chunk_count=chunk_count+${n}$ WHERE id=${kb_id}$", + {"n": ingest_result.get("chunks", 0), "kb_id": kb_id}) + except Exception as e: + exception(f"async ingest failed: {e}") + + return json.dumps({ + "status": "SUCCEEDED", + "doc_id": doc_id, + "file_name": file_name, + "file_size": file_size, + "file_type": file_type, + "ingest": ingest_result + }, ensure_ascii=False, default=str) + except Exception as e: - exception(f"ingest: {e}, {format_exc()}") + exception(f"doc_upload: {e}, {format_exc()}") return json.dumps({"error": str(e)}) +async def doc_delete_handler(request, params_kw, *args, **kwargs): + """文件删除 → 清理VDB → 清理图 → 清理DB → 删文件""" + env = request._run_ns + try: + doc_id = params_kw.get("doc_id", "") + if not doc_id: + return json.dumps({"error": "doc_id required"}) + + async with get_sor_context(env, 'rag') as sor: + # Get document info + recs = await sor.R("documents", {"id": doc_id}) + if not recs: + return json.dumps({"error": "document not found"}) + doc = recs[0] + + # Get chunks to clean VDB + chunks = await sor.R("document_chunks", {"doc_id": doc_id}) + + # Delete from VDB + if chunks: + vector_ids = [c.vector_id for c in chunks if c.vector_id] + if vector_ids: + try: + _call_vdb("/v1/delete", {"colname": doc.kb_id, "ids": vector_ids}) + except Exception as e: + exception(f"vdb delete failed: {e}") + + # Delete from graph + entities = await sor.R("entities", {"kb_id": doc.kb_id}) + if entities: + try: + _call_graph("/api/graph/save", {"graph": doc.kb_id}) + except Exception as e: + exception(f"graph cleanup failed: {e}") + + # Delete DB records + await sor.sqlExe("DELETE FROM document_chunks WHERE doc_id=${id}$", {"id": doc_id}) + await sor.sqlExe("DELETE FROM entities WHERE kb_id=${kb_id}$", {"kb_id": doc.kb_id}) + await sor.sqlExe("DELETE FROM entity_relations WHERE kb_id=${kb_id}$", {"kb_id": doc.kb_id}) + await sor.sqlExe("DELETE FROM documents WHERE id=${id}$", {"id": doc_id}) + + # Update KB stats + await sor.sqlExe( + "UPDATE knowledge_bases SET doc_count=GREATEST(doc_count-1,0), total_size=GREATEST(total_size-${size}$,0), chunk_count=GREATEST(chunk_count-${n}$,0) WHERE id=${kb_id}$", + {"size": doc.file_size, "n": len(chunks), "kb_id": doc.kb_id}) + + # Delete file from disk + file_path = doc.file_path + if file_path and file_path.startswith("/idfile/"): + real_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "files", + os.path.basename(file_path)) + if os.path.exists(real_path): + os.remove(real_path) + + return json.dumps({"status": "SUCCEEDED", "doc_id": doc_id, "chunks_deleted": len(chunks)}) + + except Exception as e: + exception(f"doc_delete: {e}, {format_exc()}") + return json.dumps({"error": str(e)}) + + +def _detect_file_type(name, mime): + """Detect file type from name and MIME""" + ext = os.path.splitext(name)[1].lower() + if ext in ('.txt', '.md', '.json', '.csv', '.xml', '.html', '.py', '.js', '.css', '.yaml', '.yml'): + return "text" + if ext in ('.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp', '.svg'): + return "image" + if ext in ('.mp3', '.wav', '.flac', '.ogg', '.m4a', '.aac'): + return "audio" + if ext in ('.mp4', '.avi', '.mov', '.mkv', '.webm'): + return "video" + if ext == '.pdf': + return "text" + return "other" + + +def _call_vdb(path, data, timeout=10): + """Call VDB service""" + import urllib.request + url = f"http://localhost:8886{path}" + req = urllib.request.Request(url, data=json.dumps(data).encode(), + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=timeout) as resp: + return json.loads(resp.read()) + + +def _call_graph(path, data, timeout=10): + """Call Graph service""" + import urllib.request + url = f"http://localhost:9092{path}" + req = urllib.request.Request(url, data=json.dumps(data).encode(), + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=timeout) as resp: + return json.loads(resp.read()) + + def init_ragserver(): - """Register API handlers""" env = ServerEnv() rf = RegisterFunction() rf.register("status", status_handler) rf.register("kb_list", kb_list_handler) rf.register("engines", engines_handler) rf.register("search", search_handler) - rf.register("ingest", ingest_handler) + rf.register("doc_upload", doc_upload_handler) + rf.register("doc_delete", doc_delete_handler)