feat: file upload → RAG ingest + file delete → cleanup pipeline
This commit is contained in:
parent
f5a91d10ec
commit
c25351a65e
@ -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
|
||||
|
||||
217
init.py
217
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)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user