111 lines
4.2 KiB
Python
111 lines
4.2 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
|
|
|
|
|
|
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"
|
|
]
|
|
}, 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})
|
|
return json.dumps({
|
|
"status": "SUCCEEDED",
|
|
"knowledge_bases": [dict(r) for r in recs]
|
|
}, 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})
|
|
return json.dumps({
|
|
"status": "SUCCEEDED",
|
|
"engines": [dict(r) for r in recs]
|
|
}, 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 ingest_handler(request, params_kw, *args, **kwargs):
|
|
"""文档入库"""
|
|
env = request._run_ns
|
|
try:
|
|
document = params_kw.get("document", "")
|
|
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"})
|
|
except Exception as e:
|
|
exception(f"ingest: {e}, {format_exc()}")
|
|
return json.dumps({"error": str(e)})
|
|
|
|
|
|
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)
|