265 lines
11 KiB
Python
265 lines
11 KiB
Python
# -*- 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})
|
|
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"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", "")
|
|
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"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():
|
|
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)
|