350 lines
14 KiB
Python
350 lines
14 KiB
Python
import json
|
||
import os
|
||
from typing import Any, Dict, List
|
||
|
||
import httpx
|
||
|
||
from config import Config
|
||
|
||
|
||
def _dump_json_content(data: Dict[str, Any]) -> str:
|
||
return json.dumps(data, ensure_ascii=False, indent=2)
|
||
|
||
|
||
def _extract_tables_map(data: Dict[str, Any]) -> Dict[str, Any]:
|
||
"""兼容两种结构：{"tables": {...}} 或直接 {...}"""
|
||
tables = data.get("tables") if isinstance(data, dict) else None
|
||
if isinstance(tables, dict):
|
||
return tables
|
||
if isinstance(data, dict):
|
||
return data
|
||
return {}
|
||
|
||
|
||
class RagflowSync:
|
||
"""RAGFlow 同步工具"""
|
||
|
||
def __init__(self):
|
||
cfg = Config.get_section("ragflow")
|
||
self._base_url = cfg.get("url", "").rstrip("/")
|
||
self._api_key = cfg.get("api_key", "")
|
||
self._table_retrieval_dataset_id = (cfg.get("table_retrieval_dataset_id") or "").strip()
|
||
self._sql_gen_dataset_id = (cfg.get("sql_gen_dataset_id") or "").strip()
|
||
|
||
def upload_documents(self, dataset_id: str, documents: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||
"""上传文档到指定知识库（每个文档单独上传）"""
|
||
if not self._base_url:
|
||
raise RuntimeError("未配置 ragflow.url")
|
||
if not dataset_id:
|
||
raise RuntimeError("dataset_id 为空，无法上传文档")
|
||
if not documents:
|
||
raise RuntimeError("没有可上传的文档内容")
|
||
|
||
# 构建正确的 URL
|
||
url = f"{self._base_url}/api/v1/datasets/{dataset_id}/documents"
|
||
|
||
headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {}
|
||
|
||
results: List[Dict[str, Any]] = []
|
||
with httpx.Client(timeout=60) as client:
|
||
for idx, doc in enumerate(documents, start=1):
|
||
content = str(doc.get("content", ""))
|
||
filename = str(doc.get("filename") or f"doc_{idx}.txt")
|
||
files = {"file": (filename, content.encode("utf-8"), "text/plain")}
|
||
|
||
print(f"上传文档 URL: {url}")
|
||
print(f"上传文件名: {filename}")
|
||
|
||
response = client.post(url, files=files, headers=headers)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
print(f"RAGFlow 上传响应: {result}")
|
||
results.append(result)
|
||
|
||
dataset_detail = self._get_dataset_detail(dataset_id)
|
||
chunk_method = self._extract_chunk_method(dataset_detail)
|
||
if chunk_method is None:
|
||
chunk_method = self._extract_chunk_method_from_upload_results(results)
|
||
# 参考 Java 实现：查询知识库文档 ID 后统一调用 chunks 解析
|
||
doc_ids = self._list_document_ids(dataset_id)
|
||
parse_results = self._auto_parse_documents(dataset_id, doc_ids)
|
||
upload_status = self._build_parse_status_from_upload_results(results)
|
||
|
||
return {
|
||
"ok": True,
|
||
"count": len(results),
|
||
"results": results,
|
||
"chunk_method": chunk_method,
|
||
"upload_status": upload_status,
|
||
"parse": parse_results,
|
||
}
|
||
|
||
def _get_dataset_detail(self, dataset_id: str) -> Dict[str, Any]:
|
||
"""查询知识库详情（用于读取 chunk_method）"""
|
||
url = f"{self._base_url}/api/v1/datasets/{dataset_id}"
|
||
headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {}
|
||
with httpx.Client(timeout=30) as client:
|
||
resp = client.get(url, headers=headers)
|
||
resp.raise_for_status()
|
||
return resp.json()
|
||
|
||
@staticmethod
|
||
def _extract_chunk_method(dataset_detail: Dict[str, Any]) -> Any:
|
||
"""从知识库详情提取 chunk_method"""
|
||
data = dataset_detail.get("data")
|
||
if isinstance(data, dict):
|
||
if "chunk_method" in data:
|
||
return data.get("chunk_method")
|
||
parser_cfg = data.get("parser_config") or {}
|
||
if isinstance(parser_cfg, dict):
|
||
return parser_cfg.get("chunk_method")
|
||
return None
|
||
|
||
@staticmethod
|
||
def _extract_chunk_method_from_upload_results(upload_results: List[Dict[str, Any]]) -> Any:
|
||
"""从上传响应中提取 chunk_method（兼容不同版本返回结构）"""
|
||
for item in upload_results:
|
||
data = item.get("data")
|
||
records = data if isinstance(data, list) else [data] if isinstance(data, dict) else []
|
||
for rec in records:
|
||
if not isinstance(rec, dict):
|
||
continue
|
||
if rec.get("chunk_method"):
|
||
return rec.get("chunk_method")
|
||
parser_cfg = rec.get("parser_config") or {}
|
||
if isinstance(parser_cfg, dict) and parser_cfg.get("chunk_method"):
|
||
return parser_cfg.get("chunk_method")
|
||
return None
|
||
|
||
@staticmethod
|
||
def _extract_uploaded_doc_ids(upload_results: List[Dict[str, Any]]) -> List[str]:
|
||
"""从上传结果中提取文档 ID"""
|
||
ids: List[str] = []
|
||
for item in upload_results:
|
||
data = item.get("data")
|
||
if isinstance(data, list):
|
||
for d in data:
|
||
if isinstance(d, dict) and d.get("id"):
|
||
ids.append(str(d.get("id")))
|
||
elif isinstance(data, dict) and data.get("id"):
|
||
ids.append(str(data.get("id")))
|
||
return ids
|
||
|
||
@staticmethod
|
||
def _build_parse_status_from_upload_results(upload_results: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||
"""根据上传返回构造解析状态（上传接口已触发解析，无需额外 parse API）"""
|
||
details: List[Dict[str, Any]] = []
|
||
for item in upload_results:
|
||
data = item.get("data")
|
||
records = data if isinstance(data, list) else [data] if isinstance(data, dict) else []
|
||
for rec in records:
|
||
if not isinstance(rec, dict):
|
||
continue
|
||
details.append(
|
||
{
|
||
"doc_id": rec.get("id"),
|
||
"name": rec.get("name") or rec.get("location"),
|
||
"run": rec.get("run"),
|
||
"chunk_method": rec.get("chunk_method")
|
||
or (rec.get("parser_config") or {}).get("chunk_method"),
|
||
}
|
||
)
|
||
return {
|
||
"ok": True,
|
||
"trigger": "upload_endpoint",
|
||
"message": "文档上传接口已触发解析流程，无需单独调用 parse API",
|
||
"count": len(details),
|
||
"details": details,
|
||
}
|
||
|
||
def _auto_parse_documents(self, dataset_id: str, doc_ids: List[str]) -> Dict[str, Any]:
|
||
"""调用官方 chunks 接口触发解析"""
|
||
if not doc_ids:
|
||
return {"ok": False, "message": "未提取到文档ID，无法触发解析", "count": 0, "details": []}
|
||
|
||
url = f"{self._base_url}/api/v1/datasets/{dataset_id}/chunks"
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"Authorization": f"Bearer {self._api_key}",
|
||
} if self._api_key else {"Content-Type": "application/json"}
|
||
payload = {"document_ids": doc_ids}
|
||
|
||
with httpx.Client(timeout=60) as client:
|
||
resp = client.post(url, headers=headers, json=payload)
|
||
|
||
if resp.status_code >= 400:
|
||
return {
|
||
"ok": False,
|
||
"trigger": "chunks_api",
|
||
"status": resp.status_code,
|
||
"message": resp.text,
|
||
"count": len(doc_ids),
|
||
"details": [{"doc_id": d} for d in doc_ids],
|
||
}
|
||
|
||
body: Any
|
||
try:
|
||
body = resp.json()
|
||
except Exception:
|
||
body = resp.text
|
||
|
||
return {
|
||
"ok": True,
|
||
"trigger": "chunks_api",
|
||
"count": len(doc_ids),
|
||
"details": [{"doc_id": d} for d in doc_ids],
|
||
"response": body,
|
||
}
|
||
|
||
def _list_document_ids(self, dataset_id: str) -> List[str]:
|
||
"""获取知识库中的全部文档 ID（用于覆盖更新）"""
|
||
if not self._base_url:
|
||
raise RuntimeError("未配置 ragflow.url")
|
||
if not dataset_id:
|
||
raise RuntimeError("dataset_id 为空，无法查询文档")
|
||
|
||
url = f"{self._base_url}/api/v1/datasets/{dataset_id}/documents"
|
||
headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {}
|
||
|
||
ids: List[str] = []
|
||
page = 1
|
||
page_size = 100
|
||
|
||
with httpx.Client(timeout=60) as client:
|
||
while True:
|
||
resp = client.get(url, headers=headers, params={"page": page, "page_size": page_size})
|
||
resp.raise_for_status()
|
||
body = resp.json()
|
||
data = body.get("data")
|
||
if isinstance(data, dict):
|
||
docs = data.get("docs") or data.get("list") or []
|
||
elif isinstance(data, list):
|
||
docs = data
|
||
else:
|
||
docs = []
|
||
|
||
if not docs:
|
||
break
|
||
|
||
for item in docs:
|
||
if isinstance(item, dict) and item.get("id"):
|
||
ids.append(str(item.get("id")))
|
||
|
||
if len(docs) < page_size:
|
||
break
|
||
page += 1
|
||
|
||
return ids
|
||
|
||
def _delete_documents(self, dataset_id: str, doc_ids: List[str]) -> Dict[str, Any]:
|
||
"""按 ID 删除文档"""
|
||
if not doc_ids:
|
||
return {"ok": True, "deleted": 0}
|
||
|
||
url = f"{self._base_url}/api/v1/datasets/{dataset_id}/documents"
|
||
headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {}
|
||
payload = {"ids": doc_ids}
|
||
|
||
with httpx.Client(timeout=60) as client:
|
||
resp = client.request("DELETE", url, headers=headers, json=payload)
|
||
resp.raise_for_status()
|
||
return resp.json()
|
||
|
||
def replace_documents(self, dataset_id: str, documents: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||
"""覆盖更新：先删后传，避免“update 变新增”"""
|
||
ids = self._list_document_ids(dataset_id)
|
||
if ids:
|
||
self._delete_documents(dataset_id, ids)
|
||
return self.upload_documents(dataset_id, documents)
|
||
|
||
def update_table_retrieval_documents(self) -> Dict[str, Any]:
|
||
"""更新表名检索文档（仅文档内容）"""
|
||
if not self._table_retrieval_dataset_id:
|
||
raise RuntimeError("未配置 ragflow.table_retrieval_dataset_id，无法更新表名检索文档")
|
||
|
||
root = os.path.dirname(os.path.dirname(__file__))
|
||
tables_file = os.path.join(root, "config", "table_retrieval_prompts", "tables.json")
|
||
with open(tables_file, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
|
||
tables = _extract_tables_map(data)
|
||
documents = [
|
||
{
|
||
"filename": f"{k}.txt",
|
||
"content": _dump_json_content({"table": k, "templates": v}),
|
||
}
|
||
for k, v in tables.items()
|
||
]
|
||
return self.replace_documents(self._table_retrieval_dataset_id, documents)
|
||
|
||
def update_sql_gen_documents(self) -> Dict[str, Any]:
|
||
"""更新 SQL 生成文档（仅文档内容）"""
|
||
if not self._sql_gen_dataset_id:
|
||
raise RuntimeError("未配置 ragflow.sql_gen_dataset_id，无法更新 SQL 生成文档")
|
||
|
||
root = os.path.dirname(os.path.dirname(__file__))
|
||
prompts_dir = os.path.join(root, "config", "sql_gen_prompts")
|
||
documents: List[Dict[str, Any]] = []
|
||
|
||
for name in os.listdir(prompts_dir):
|
||
if not name.endswith(".json"):
|
||
continue
|
||
path = os.path.join(prompts_dir, name)
|
||
with open(path, "r", encoding="utf-8") as f:
|
||
prompt = json.load(f)
|
||
table = prompt.get("table") or os.path.splitext(name)[0]
|
||
documents.append({"filename": f"{table}.txt", "content": _dump_json_content(prompt)})
|
||
|
||
return self.replace_documents(self._sql_gen_dataset_id, documents)
|
||
|
||
def upload_table_retrieval(self) -> Dict[str, Any]:
|
||
"""上传表名检索模板文档 - 直接上传整个 JSON 文件"""
|
||
if not self._table_retrieval_dataset_id:
|
||
raise RuntimeError("未配置 ragflow.table_retrieval_dataset_id，无法上传表名检索模板")
|
||
|
||
root = os.path.dirname(os.path.dirname(__file__))
|
||
tables_file = os.path.join(root, "config", "table_retrieval_prompts", "tables.json")
|
||
|
||
if not os.path.exists(tables_file):
|
||
raise RuntimeError(f"表名检索模板文件不存在: {tables_file}")
|
||
|
||
# 读取整个 JSON 文件内容
|
||
with open(tables_file, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
|
||
tables = _extract_tables_map(data)
|
||
# 每个 key 一个文档，配合 One 解析时每个表单独成块
|
||
documents = [
|
||
{
|
||
"filename": f"{k}.txt",
|
||
"content": _dump_json_content({"table": k, "templates": v}),
|
||
}
|
||
for k, v in tables.items()
|
||
]
|
||
return self.upload_documents(self._table_retrieval_dataset_id, documents)
|
||
|
||
def upload_sql_gen(self) -> Dict[str, Any]:
|
||
"""上传 SQL 生成提示词文档"""
|
||
if not self._sql_gen_dataset_id:
|
||
raise RuntimeError("未配置 ragflow.sql_gen_dataset_id，无法上传 SQL 生成提示词")
|
||
|
||
root = os.path.dirname(os.path.dirname(__file__))
|
||
prompts_dir = os.path.join(root, "config", "sql_gen_prompts")
|
||
|
||
if not os.path.exists(prompts_dir):
|
||
raise RuntimeError(f"SQL 生成提示词目录不存在: {prompts_dir}")
|
||
|
||
documents = []
|
||
|
||
for name in os.listdir(prompts_dir):
|
||
if not name.endswith(".json"):
|
||
continue
|
||
path = os.path.join(prompts_dir, name)
|
||
with open(path, "r", encoding="utf-8") as f:
|
||
prompt = json.load(f)
|
||
table = prompt.get("table") or os.path.splitext(name)[0]
|
||
json_content = _dump_json_content(prompt)
|
||
documents.append({"filename": f"{table}.txt", "content": json_content})
|
||
|
||
return self.upload_documents(self._sql_gen_dataset_id, documents)
|