298 lines
12 KiB
Python
298 lines
12 KiB
Python
"""知识库服务：调用 BaoDan Dataset API 管理文档。"""
|
|
import requests
|
|
from flask import current_app
|
|
|
|
|
|
class KBService:
|
|
"""知识库管理业务逻辑。"""
|
|
|
|
def _baodan_request(self, method: str, path: str, **kwargs):
|
|
"""发送 BaoDan API 请求的通用方法。"""
|
|
api_key = current_app.config.get("BAODAN_KB_API_KEY", "")
|
|
base_url = current_app.config.get("BAODAN_API_URL", "http://localhost:5001")
|
|
headers = {"Authorization": f"Bearer {api_key}"}
|
|
try:
|
|
resp = requests.request(method, f"{base_url}{path}", headers=headers, timeout=30, **kwargs)
|
|
return resp.json()
|
|
except requests.Timeout:
|
|
return {"code": 2002, "message": "BaoDan API 超时", "data": None}
|
|
except Exception as e:
|
|
return {"code": 2001, "message": f"BaoDan API 调用失败: {str(e)}", "data": None}
|
|
|
|
def upload_documents(self, files, insurance_type: str, company: str) -> dict:
|
|
"""上传文档到 BaoDan 知识库。"""
|
|
api_key = current_app.config.get("BAODAN_KB_API_KEY", "")
|
|
base_url = current_app.config.get("BAODAN_API_URL", "http://localhost:5001")
|
|
|
|
uploaded = []
|
|
for f in files:
|
|
resp = requests.post(
|
|
f"{base_url}/v1/datasets/documents/upload",
|
|
headers={"Authorization": f"Bearer {api_key}"},
|
|
files={"file": (f.filename, f.stream, f.content_type)},
|
|
data={
|
|
"indexing_technique": "high_quality",
|
|
"process_rule_mode": "automatic",
|
|
},
|
|
timeout=60,
|
|
)
|
|
data = resp.json()
|
|
uploaded.append({
|
|
"doc_id": data.get("id", ""),
|
|
"filename": f.filename,
|
|
"status": "processing",
|
|
})
|
|
|
|
# 记录操作日志
|
|
from insurance.utils.audit import log_operation
|
|
log_operation("system", "upload", "document", "", {
|
|
"files": [f.filename for f in files],
|
|
"insurance_type": insurance_type,
|
|
"company": company,
|
|
})
|
|
|
|
return {"code": 0, "data": {"upload_id": "", "documents": uploaded}}
|
|
|
|
def list_documents(self, params: dict) -> dict:
|
|
"""获取文档列表。"""
|
|
result = self._baodan_request("GET", "/v1/datasets/documents", params={
|
|
"page": params["page"],
|
|
"limit": params["page_size"],
|
|
})
|
|
|
|
# BaoDan API 返回 {code: 0, data: [...]}，直接透传
|
|
if result and "data" in result:
|
|
documents = result["data"]
|
|
if params.get("keyword"):
|
|
keyword = params["keyword"].lower()
|
|
documents = [d for d in documents if keyword in d.get("name", "").lower()]
|
|
return {"code": 0, "data": documents}
|
|
|
|
return {"code": 0, "data": []}
|
|
|
|
def get_document_status(self, doc_id: str) -> dict:
|
|
"""查询文档处理状态。"""
|
|
result = self._baodan_request("GET", f"/v1/datasets/documents/{doc_id}/status")
|
|
return {"code": 0, "data": result}
|
|
|
|
def retry_document(self, doc_id: str) -> dict:
|
|
"""重试失败文档。"""
|
|
result = self._baodan_request("POST", f"/v1/datasets/documents/{doc_id}/retry")
|
|
return {"code": 0, "data": result}
|
|
|
|
def update_document(self, doc_id: str, data: dict) -> dict:
|
|
"""更新文档元数据。"""
|
|
result = self._baodan_request("PATCH", f"/v1/datasets/documents/{doc_id}", json=data)
|
|
return {"code": 0, "data": result}
|
|
|
|
def delete_document(self, doc_id: str) -> dict:
|
|
"""删除文档。"""
|
|
result = self._baodan_request("DELETE", f"/v1/datasets/documents/{doc_id}")
|
|
|
|
# 记录操作日志
|
|
from insurance.utils.audit import log_operation
|
|
log_operation("system", "delete", "document", doc_id)
|
|
|
|
return {"code": 0, "message": "success", "data": None}
|
|
|
|
def generate_doc_number(self, insurance_type: str, company: str) -> str:
|
|
"""生成文档编号。"""
|
|
from insurance.db.compat import db
|
|
from sqlalchemy import text
|
|
|
|
type_codes = {
|
|
"重疾险": "CJ", "寿险": "RS", "医疗险": "YL",
|
|
"意外险": "YW", "车险": "CX", "年金险": "NJ", "储蓄险": "CX",
|
|
}
|
|
type_code = type_codes.get(insurance_type, "XX")
|
|
company_code = company[:2] if company else "XX"
|
|
|
|
# 查询当前序号
|
|
result = db.session.execute(
|
|
text("""
|
|
SELECT COUNT(*) FROM document_metadata
|
|
WHERE insurance_type = :type AND company = :company
|
|
"""),
|
|
{"type": insurance_type, "company": company}
|
|
)
|
|
count = result.scalar() or 0
|
|
|
|
return f"{type_code}-{company_code}-{count + 1:03d}"
|
|
|
|
def list_datasources(self) -> dict:
|
|
"""获取保司 API 数据源列表。"""
|
|
from insurance.db.compat import db
|
|
from sqlalchemy import text
|
|
try:
|
|
result = db.session.execute(text("SELECT * FROM kb_datasources ORDER BY created_at DESC"))
|
|
return {"code": 0, "data": [dict(row) for row in result]}
|
|
except Exception:
|
|
return {"code": 0, "data": []}
|
|
|
|
def create_datasource(self, data: dict) -> dict:
|
|
"""新建数据源配置。"""
|
|
from insurance.db.compat import db
|
|
from sqlalchemy import text
|
|
try:
|
|
db.session.execute(
|
|
text("INSERT INTO kb_datasources (name, api_url, api_key, sync_frequency) VALUES (:name, :url, :key, :freq)"),
|
|
{"name": data["name"], "url": data["api_url"], "key": data.get("api_key", ""), "freq": data.get("sync_frequency", "")},
|
|
)
|
|
db.session.commit()
|
|
|
|
# 记录操作日志
|
|
from insurance.utils.audit import log_operation
|
|
log_operation("system", "create", "datasource", "", {"name": data["name"]})
|
|
|
|
return {"code": 0, "message": "success", "data": None}
|
|
except Exception as e:
|
|
return {"code": 5001, "message": f"创建数据源失败: {str(e)}", "data": None}
|
|
|
|
def sync_datasource(self, ds_id: str) -> dict:
|
|
"""手动触发同步。"""
|
|
import threading
|
|
from insurance.db.compat import db
|
|
from sqlalchemy import text
|
|
|
|
# 获取数据源配置
|
|
try:
|
|
result = db.session.execute(
|
|
text("SELECT * FROM kb_datasources WHERE id = :ds_id"),
|
|
{"ds_id": ds_id},
|
|
)
|
|
datasource = result.fetchone()
|
|
if not datasource:
|
|
return {"code": 1005, "message": "数据源不存在", "data": None}
|
|
except Exception as e:
|
|
return {"code": 5001, "message": f"查询数据源失败: {str(e)}", "data": None}
|
|
|
|
# 记录同步开始
|
|
self._log_sync_start(ds_id)
|
|
|
|
# 异步执行同步任务
|
|
from flask import current_app
|
|
flask_app = current_app._get_current_object()
|
|
|
|
def sync_task():
|
|
with flask_app.app_context():
|
|
try:
|
|
self._execute_sync(ds_id, dict(datasource) if datasource else {})
|
|
self._log_sync_success(ds_id)
|
|
except Exception as e:
|
|
self._log_sync_failure(ds_id, str(e))
|
|
from insurance.utils.notify import notify_sync_failure
|
|
notify_sync_failure(f"数据源#{ds_id}", str(e))
|
|
|
|
thread = threading.Thread(target=sync_task, daemon=True)
|
|
thread.start()
|
|
|
|
return {"code": 0, "message": "同步任务已提交", "data": None}
|
|
|
|
def _log_sync_start(self, ds_id: str):
|
|
"""记录同步开始。"""
|
|
from insurance.db.compat import db
|
|
from sqlalchemy import text
|
|
try:
|
|
db.session.execute(
|
|
text("INSERT INTO kb_sync_logs (datasource_id, status, message) VALUES (:ds_id, :status, :message)"),
|
|
{"ds_id": ds_id, "status": "running", "message": "同步开始"},
|
|
)
|
|
db.session.commit()
|
|
except Exception:
|
|
db.session.rollback()
|
|
|
|
def _log_sync_success(self, ds_id: str):
|
|
"""记录同步成功。"""
|
|
from insurance.db.compat import db
|
|
from sqlalchemy import text
|
|
try:
|
|
db.session.execute(
|
|
text("UPDATE kb_sync_logs SET status = :status, message = :message WHERE datasource_id = :ds_id AND status = 'running'"),
|
|
{"ds_id": ds_id, "status": "success", "message": "同步完成"},
|
|
)
|
|
db.session.commit()
|
|
except Exception:
|
|
db.session.rollback()
|
|
|
|
def _log_sync_failure(self, ds_id: str, error: str):
|
|
"""记录同步失败。"""
|
|
from insurance.db.compat import db
|
|
from sqlalchemy import text
|
|
try:
|
|
db.session.execute(
|
|
text("UPDATE kb_sync_logs SET status = :status, message = :message WHERE datasource_id = :ds_id AND status = 'running'"),
|
|
{"ds_id": ds_id, "status": "failed", "message": error[:500]},
|
|
)
|
|
db.session.commit()
|
|
except Exception:
|
|
db.session.rollback()
|
|
|
|
def _execute_sync(self, ds_id: str, datasource: dict):
|
|
"""执行实际的同步操作。"""
|
|
import requests
|
|
from flask import current_app
|
|
|
|
api_url = datasource.get("api_url", "")
|
|
api_key = datasource.get("api_key", "")
|
|
|
|
if not api_url:
|
|
raise Exception("API地址未配置")
|
|
|
|
# 调用保司API获取数据
|
|
headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
|
|
resp = requests.get(api_url, headers=headers, timeout=60)
|
|
resp.raise_for_status()
|
|
|
|
data = resp.json()
|
|
# TODO: 根据实际API响应格式解析数据并同步到知识库
|
|
# 这里需要根据具体保司API格式进行适配
|
|
|
|
return data
|
|
|
|
def get_sync_logs(self, ds_id: str) -> dict:
|
|
"""获取同步日志。"""
|
|
from insurance.db.compat import db
|
|
from sqlalchemy import text
|
|
try:
|
|
result = db.session.execute(
|
|
text("SELECT * FROM kb_sync_logs WHERE datasource_id = :ds_id ORDER BY created_at DESC LIMIT 50"),
|
|
{"ds_id": ds_id},
|
|
)
|
|
return {"code": 0, "data": [dict(row) for row in result]}
|
|
except Exception:
|
|
return {"code": 0, "data": []}
|
|
|
|
def retrieval_search(self, query: str, filters: dict, top_k: int) -> dict:
|
|
"""向量检索：调用 BaoDan Dataset Retrieval API。"""
|
|
api_key = current_app.config.get("BAODAN_KB_API_KEY", "")
|
|
base_url = current_app.config.get("BAODAN_API_URL", "http://localhost:5001")
|
|
|
|
try:
|
|
resp = requests.post(
|
|
f"{base_url}/v1/datasets/retrieve",
|
|
json={
|
|
"query": query,
|
|
"retrieval_model": {
|
|
"search_method": "semantic_search",
|
|
"reranking_enable": False,
|
|
"top_k": top_k,
|
|
"score_threshold_enabled": False,
|
|
},
|
|
},
|
|
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
|
|
timeout=30,
|
|
)
|
|
data = resp.json()
|
|
results = [
|
|
{
|
|
"doc_name": r.get("document_name", ""),
|
|
"chunk": r.get("segment_content", ""),
|
|
"score": r.get("score", 0),
|
|
"metadata": r.get("metadata", {}),
|
|
}
|
|
for r in data.get("records", [])
|
|
]
|
|
return {"code": 0, "data": {"results": results}}
|
|
except Exception as e:
|
|
return {"code": 2001, "message": f"检索失败: {str(e)}", "data": None}
|