init
This commit is contained in:
+273
-183
@@ -7,45 +7,18 @@ import httpx
|
||||
from config import Config
|
||||
|
||||
|
||||
def _build_document_for_table(table: str, templates: List[str]) -> str:
|
||||
"""构建表名检索文档 - 使用更标准的格式"""
|
||||
lines = [
|
||||
f"# 表名检索模板: {table}",
|
||||
"",
|
||||
"## 可用模板:",
|
||||
""
|
||||
]
|
||||
for i, t in enumerate(templates, 1):
|
||||
lines.append(f"{i}. {t}")
|
||||
lines.extend(["", f"表名: {table}", "类型: 表名检索模板"])
|
||||
return "\n".join(lines)
|
||||
def _dump_json_content(data: Dict[str, Any]) -> str:
|
||||
return json.dumps(data, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
def _build_sql_gen_document(table: str, prompt: Dict[str, any]) -> str:
|
||||
"""构建 SQL 生成文档 - 使用更标准的格式"""
|
||||
system_prompt = prompt.get("system_prompt", "")
|
||||
business_prompt = prompt.get("business_prompt", "")
|
||||
constraints = prompt.get("constraints", [])
|
||||
|
||||
lines = [
|
||||
f"# SQL 生成提示词: {table}",
|
||||
"",
|
||||
"## 系统提示词:",
|
||||
system_prompt,
|
||||
"",
|
||||
"## 业务提示词:",
|
||||
business_prompt,
|
||||
""
|
||||
]
|
||||
|
||||
if constraints:
|
||||
lines.extend(["## 约束条件:", ""])
|
||||
for i, c in enumerate(constraints, 1):
|
||||
lines.append(f"{i}. {c}")
|
||||
lines.append("")
|
||||
|
||||
lines.extend([f"表名: {table}", "类型: SQL 生成提示词"])
|
||||
return "\n".join(lines)
|
||||
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:
|
||||
@@ -58,143 +31,271 @@ class RagflowSync:
|
||||
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 _validate_common(self) -> None:
|
||||
if not self._base_url:
|
||||
raise RuntimeError("未配置 ragflow.url")
|
||||
if not self._upload_path:
|
||||
raise RuntimeError("未配置 ragflow.upload 上传接口，请在 config/config.ini 中设置")
|
||||
if "{dataset_id}" not in self._upload_path:
|
||||
raise RuntimeError("上传接口路径必须包含 {dataset_id} 占位符")
|
||||
if self._upload_mode not in ("overwrite", "append"):
|
||||
raise RuntimeError("ragflow.upload_mode 仅支持 overwrite 或 append")
|
||||
|
||||
def _post(self, documents: List[Dict[str, Any]], dataset_id: str):
|
||||
"""上传文档到指定知识库 - 使用 multipart/form-data 格式"""
|
||||
if not self._base_url:
|
||||
raise RuntimeError("未配置 ragflow.url")
|
||||
|
||||
# 构建正确的 URL
|
||||
upload_path = self._upload_path.replace("{dataset_id}", dataset_id)
|
||||
url = self._base_url + "/" + upload_path.lstrip("/")
|
||||
|
||||
headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {}
|
||||
|
||||
# 由于接口使用 multipart/form-data，我们需要创建临时文件
|
||||
import tempfile
|
||||
|
||||
# 创建临时文件并写入文档内容
|
||||
with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False, encoding='utf-8') as f:
|
||||
# 将文档内容写入文件
|
||||
for doc in documents:
|
||||
content = doc.get('content', '')
|
||||
f.write(content + '\n\n')
|
||||
temp_file_path = f.name
|
||||
|
||||
try:
|
||||
# 使用 multipart/form-data 上传文件
|
||||
files = {'file': open(temp_file_path, 'rb')}
|
||||
|
||||
print(f"请求 URL: {url}") # 调试信息
|
||||
print(f"上传文件: {temp_file_path}") # 调试信息
|
||||
|
||||
with httpx.Client(timeout=60) as client:
|
||||
response = client.post(url, files=files, headers=headers)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
print(f"RAGFlow 上传响应: {result}") # 调试信息
|
||||
return result
|
||||
finally:
|
||||
# 清理临时文件
|
||||
import os
|
||||
if os.path.exists(temp_file_path):
|
||||
os.unlink(temp_file_path)
|
||||
|
||||
def upload_documents(self, dataset_id: str, documents: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||||
"""上传文档到指定知识库 - 使用 multipart/form-data 格式
|
||||
|
||||
根据官方文档: POST /api/v1/datasets/{dataset_id}/documents
|
||||
"""
|
||||
"""上传文档到指定知识库（每个文档单独上传）"""
|
||||
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 {}
|
||||
|
||||
# 由于接口使用 multipart/form-data，我们需要创建临时文件
|
||||
import tempfile
|
||||
import os
|
||||
|
||||
# 创建临时文件并写入文档内容
|
||||
with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False, encoding='utf-8') as f:
|
||||
# 将文档内容写入文件
|
||||
for doc in documents:
|
||||
content = doc.get('content', '')
|
||||
f.write(content + '\n\n')
|
||||
temp_file_path = f.name
|
||||
|
||||
try:
|
||||
# 使用 multipart/form-data 上传文件
|
||||
# 确保文件在 with 块内打开和关闭
|
||||
with open(temp_file_path, 'rb') as file_obj:
|
||||
files = {'file': file_obj}
|
||||
|
||||
print(f"上传文档 URL: {url}") # 调试信息
|
||||
print(f"上传文件: {temp_file_path}") # 调试信息
|
||||
|
||||
with httpx.Client(timeout=60) as client:
|
||||
response = client.post(url, files=files, headers=headers)
|
||||
|
||||
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}") # 调试信息
|
||||
|
||||
# 检查文档处理状态
|
||||
if result.get('code') == 0 and result.get('data'):
|
||||
doc_id = result['data'][0].get('id')
|
||||
if doc_id:
|
||||
print(f"文档已上传，ID: {doc_id}")
|
||||
print("注意: 文档处理需要时间，请等待 RAGFlow 完成分块处理")
|
||||
print("可以在 RAGFlow 界面查看处理进度")
|
||||
|
||||
return result
|
||||
finally:
|
||||
# 清理临时文件
|
||||
if os.path.exists(temp_file_path):
|
||||
try:
|
||||
os.unlink(temp_file_path)
|
||||
except PermissionError:
|
||||
# 如果文件被占用，等待一下再重试
|
||||
import time
|
||||
time.sleep(0.1)
|
||||
try:
|
||||
os.unlink(temp_file_path)
|
||||
except PermissionError:
|
||||
print(f"警告: 无法删除临时文件 {temp_file_path}")
|
||||
print(f"RAGFlow 上传响应: {result}")
|
||||
results.append(result)
|
||||
|
||||
def update_dataset(self, dataset_id: str, config: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""更新知识库配置
|
||||
|
||||
根据官方文档: PUT /api/v1/datasets/{dataset_id}
|
||||
"""
|
||||
if not self._base_url:
|
||||
raise RuntimeError("未配置 ragflow.url")
|
||||
|
||||
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 ""
|
||||
}
|
||||
|
||||
print(f"更新知识库 URL: {url}") # 调试信息
|
||||
print(f"更新配置: {config}") # 调试信息
|
||||
|
||||
"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:
|
||||
response = client.put(url, json=config, headers=headers)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
print(f"RAGFlow 更新响应: {result}") # 调试信息
|
||||
return result
|
||||
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 文件"""
|
||||
@@ -210,22 +311,17 @@ class RagflowSync:
|
||||
# 读取整个 JSON 文件内容
|
||||
with open(tables_file, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
|
||||
# 将 JSON 内容转换为字符串
|
||||
json_content = json.dumps(data, ensure_ascii=False, indent=2)
|
||||
|
||||
# 构建文档
|
||||
doc = {
|
||||
"content": f"# 表名检索模板库\n\n以下是所有表名检索模板的 JSON 数据：\n\n```json\n{json_content}\n```\n\n包含的表：{list(data.keys())}",
|
||||
"metadata": {"type": "table_retrieval_templates", "format": "json"},
|
||||
"title": "表名检索模板库",
|
||||
"type": "table_template_library"
|
||||
}
|
||||
|
||||
print(f"生成的表名检索文档: {doc}")
|
||||
|
||||
# 上传整个 JSON 文件内容
|
||||
return self.upload_documents(self._table_retrieval_dataset_id, [doc])
|
||||
|
||||
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 生成提示词文档"""
|
||||
@@ -247,13 +343,7 @@ class RagflowSync:
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
prompt = json.load(f)
|
||||
table = prompt.get("table") or os.path.splitext(name)[0]
|
||||
doc = {
|
||||
"content": _build_sql_gen_document(table, prompt),
|
||||
"metadata": {"table": table},
|
||||
"title": f"SQL Prompt: {table}",
|
||||
"type": "sql_prompt"
|
||||
}
|
||||
documents.append(doc)
|
||||
print(f"生成的 SQL 提示词文档: {doc}")
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user