# -*- coding:utf-8 -*- """ 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, 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/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}) cards = [] for r in recs: doc_recs = await sor.R("documents", {"kb_id": r.id}) doc_count = len(doc_recs) total_size = r.total_size or 0 size_str = f"{total_size/1024:.0f}KB" if total_size else "0B" engine = r.embedding_engine or "CLIP" card = {"widgettype":"VBox","options":{"cwidth":16,"cheight":12,"bgcolor":"#f0f7ff","padding":"16px","css":"card clickable","border":"1px solid #d0e4f7"},"subwidgets":[ {"widgettype":"Text","options":{"text":"📚 " + str(r.name),"cfontsize":16,"fontWeight":"bold"}}, {"widgettype":"Text","options":{"text":str(engine)+" · 1024维","cfontsize":12,"color":"#888"}}, {"widgettype":"HBox","options":{"spacing":"12px"},"subwidgets":[ {"widgettype":"Text","options":{"text":"📄 "+str(doc_count)+"文档","cfontsize":12,"color":"#666"}}, {"widgettype":"Text","options":{"text":"💾 "+size_str,"cfontsize":12,"color":"#666"}}]}], "binds":[{"wid":"self","event":"click","actiontype":"urlwidget","target":"app.rag_main_content","mode":"replace","options":{"url":"/knowledge_bases_list/detail.ui"}}]} cards.append(card) if not cards: cards.append({"widgettype":"Text","options":{"text":"暂无知识库","color":"#aaa","cfontsize":14,"halign":"center"}}) result = {"widgettype":"HBox","options":{"width":"100%","spacing":"12px","wrap":True},"subwidgets":cards} return json.dumps(result, ensure_ascii=False) 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() async with get_sor_context(env, 'rag') as sor: sql = "SELECT * FROM engine_configs WHERE status='active' AND (org_id IS NULL OR org_id=${org_id}$) ORDER BY engine_type, priority DESC" recs = await sor.sqlExe(sql, {"org_id": userorgid}) rows = [dict(r) for r in recs] return json.dumps({"status": "SUCCEEDED", "data": {"rows": rows, "total": len(rows)}}, ensure_ascii=False, default=str) except Exception as e: exception(f"engines: {e}, {format_exc()}") return json.dumps({"error": str(e)}) async def search_handler(request, params_kw, *args, **kwargs): env = request._run_ns try: query = params_kw.get("query", "") kb_id = params_kw.get("kb_id", "") top_k = int(params_kw.get("top_k", 5)) if not query: return json.dumps({"error": "query required"}) try: from pipeline import search as pipeline_search result = pipeline_search(query, pipeline_name="kg-rag-standard", collection=kb_id or "knowledge", graph_name=kb_id or "knowledge", top_k=top_k, llm_func=None) return json.dumps(result, ensure_ascii=False) except ImportError: return json.dumps({"status": "FALLBACK", "message": "pipeline not available"}) except Exception as e: exception(f"search: {e}, {format_exc()}") return json.dumps({"error": str(e)}) async def doc_upload_handler(request, params_kw, *args, **kwargs): """文件上传 → 保存 → DB记录 → 触发RAG入库""" env = request._run_ns try: userorgid = await env.get_userorgid() kb_id = params_kw.get("kb_id", "") file_name = params_kw.get("file_name", "upload.bin") if not kb_id: return json.dumps({"error": "kb_id required"}) # Read raw file data from request body file_data = await request.read() if not file_data: return json.dumps({"error": "no file data"}) # 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, "application/octet-stream") # 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": "application/octet-stream", "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"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()) async def dir_create_handler(request, params_kw, *args, **kwargs): """创建目录""" env = request._run_ns try: kb_id = params_kw.get("kb_id", "") parent_id = params_kw.get("parent_id", "") dir_name = params_kw.get("dir_name", "") if not kb_id or not dir_name: return json.dumps({"error": "kb_id and dir_name required"}) async with get_sor_context(env, 'rag') as sor: dir_id = uuid.uuid4().hex[:16] await sor.sqlExe( "INSERT INTO document_chunks (id, doc_id, kb_id, chunk_index, chunk_type, content, description, created_at) " "VALUES (${id}$, '', ${kb_id}$, 0, 'directory', ${name}$, ${parent}$, NOW())", {"id": dir_id, "kb_id": kb_id, "name": dir_name, "parent": parent_id}) return json.dumps({"status": "SUCCEEDED", "dir_id": dir_id}) except Exception as e: exception(f"dir_create: {e}") return json.dumps({"error": str(e)}) async def dir_delete_handler(request, params_kw, *args, **kwargs): """删除目录/文件""" env = request._run_ns try: item_id = params_kw.get("item_id", "") if not item_id: return json.dumps({"error": "item_id required"}) async with get_sor_context(env, 'rag') as sor: await sor.sqlExe("DELETE FROM document_chunks WHERE id=${id}$ OR doc_id=${id}$", {"id": item_id}) return json.dumps({"status": "SUCCEEDED"}) except Exception as e: exception(f"dir_delete: {e}") return json.dumps({"error": str(e)}) async def dir_list_handler(request, params_kw, *args, **kwargs): """列出知识库的目录树""" env = request._run_ns try: kb_id = params_kw.get("kb_id", "") if not kb_id: return json.dumps({"error": "kb_id required"}) async with get_sor_context(env, 'rag') as sor: # Get directories (chunk_type='directory') dirs = await sor.sqlExe( "SELECT id, content as label, description as parent_id FROM document_chunks WHERE kb_id=${kb_id}$ AND chunk_type='directory'", {"kb_id": kb_id}) # Get files (documents table) docs = await sor.sqlExe( "SELECT id, file_name as label, '' as parent_id FROM documents WHERE kb_id=${kb_id}$", {"kb_id": kb_id}) items = [] for d in dirs: items.append({"id": d.id, "label": d.label, "parent_id": d.parent_id or "", "type": "dir"}) for d in docs: items.append({"id": d.id, "label": d.label, "parent_id": "", "type": "file"}) return json.dumps({"status": "SUCCEEDED", "items": items}) except Exception as e: exception(f"dir_list: {e}") return json.dumps({"error": str(e)}) def init_ragserver(): 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("doc_upload", doc_upload_handler) rf.register("doc_delete", doc_delete_handler) rf.register("dir_create", dir_create_handler) rf.register("dir_delete", dir_delete_handler) rf.register("dir_list", dir_list_handler)