This commit is contained in:
2026-02-26 19:23:54 +08:00
parent ad0082d811
commit a158cbbe9c
3 changed files with 260 additions and 69 deletions
+42 -5
View File
@@ -112,12 +112,49 @@ def reload_prompts(prompt_manager=Depends(get_prompt_manager)):
@router.post("/api/ragflow/table-retrieval/reload") @router.post("/api/ragflow/table-retrieval/reload")
def reload_table_retrieval(): def reload_table_retrieval():
syncer = RagflowSync() syncer = RagflowSync()
result = syncer.sync_table_retrieval() result = syncer.upload_table_retrieval()
return {"ok": True, "result": result} return {"ok": True, "result": result}
@router.post("/api/ragflow/sql-gen/reload") @router.post("/api/ragflow/table-retrieval/upload")
def reload_sql_gen(): def upload_table_retrieval():
"""上传表名检索模板文档"""
syncer = RagflowSync() syncer = RagflowSync()
result = syncer.sync_sql_gen_prompts() try:
return {"ok": True, "result": result} result = syncer.upload_table_retrieval()
return {"ok": True, "result": result}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@router.put("/api/ragflow/table-retrieval/update")
def update_table_retrieval(config: dict):
"""更新表名检索知识库配置"""
syncer = RagflowSync()
try:
result = syncer.update_dataset(syncer._table_retrieval_dataset_id, config)
return {"ok": True, "result": result}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@router.post("/api/ragflow/sql-gen/upload")
def upload_sql_gen():
"""上传 SQL 生成提示词文档"""
syncer = RagflowSync()
try:
result = syncer.upload_sql_gen()
return {"ok": True, "result": result}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@router.put("/api/ragflow/sql-gen/update")
def update_sql_gen(config: dict):
"""更新 SQL 生成知识库配置"""
syncer = RagflowSync()
try:
result = syncer.update_dataset(syncer._sql_gen_dataset_id, config)
return {"ok": True, "result": result}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
+23 -22
View File
@@ -1,23 +1,24 @@
{ {
"tables": { "apbo_eta_ful": [
"apbo_eta_ful": [ "SO ETA INFO",
"SO ETA INFO", "BO list",
"BO list", "CC",
"CC", "Commodity Code",
"Commodity Code", "BO QTY",
"BO QTY", "open order",
"open order", "BO status",
"BO status", "BO report",
"BO report", "model",
"model", "MTM SN",
"MTM SN", "backlog",
"backlog", "Premier",
"Premier", "SN machine_sn",
"SN machine_sn", "work order information",
"work order information", "recovery ETA",
"recovery ETA", "history order",
"history order", "Warranty type"
"Warranty type" ],
] "x_table_name": [
} "xx"
} ]
}
+195 -42
View File
@@ -1,6 +1,6 @@
import json import json
import os import os
from typing import Dict, List, Optional from typing import Any, Dict, List
import httpx import httpx
@@ -8,25 +8,43 @@ from config import Config
def _build_document_for_table(table: str, templates: List[str]) -> str: def _build_document_for_table(table: str, templates: List[str]) -> str:
lines = [f"table: {table}"] """构建表名检索文档 - 使用更标准的格式"""
for t in templates: lines = [
lines.append(f"- {t}") f"# 表名检索模板: {table}",
"",
"## 可用模板:",
""
]
for i, t in enumerate(templates, 1):
lines.append(f"{i}. {t}")
lines.extend(["", f"表名: {table}", "类型: 表名检索模板"])
return "\n".join(lines) return "\n".join(lines)
def _build_sql_gen_document(table: str, prompt: Dict[str, any]) -> str: def _build_sql_gen_document(table: str, prompt: Dict[str, any]) -> str:
"""构建 SQL 生成文档 - 使用更标准的格式"""
system_prompt = prompt.get("system_prompt", "") system_prompt = prompt.get("system_prompt", "")
business_prompt = prompt.get("business_prompt", "") business_prompt = prompt.get("business_prompt", "")
constraints = prompt.get("constraints", []) constraints = prompt.get("constraints", [])
lines = [f"table: {table}"]
if system_prompt: lines = [
lines.append("[system] " + system_prompt) f"# SQL 生成提示词: {table}",
if business_prompt: "",
lines.append("[business] " + business_prompt) "## 系统提示词:",
system_prompt,
"",
"## 业务提示词:",
business_prompt,
""
]
if constraints: if constraints:
lines.append("[constraints]") lines.extend(["## 约束条件:", ""])
for c in constraints: for i, c in enumerate(constraints, 1):
lines.append(f"- {c}") lines.append(f"{i}. {c}")
lines.append("")
lines.extend([f"表名: {table}", "类型: SQL 生成提示词"])
return "\n".join(lines) return "\n".join(lines)
@@ -39,54 +57,187 @@ class RagflowSync:
self._api_key = cfg.get("api_key", "") self._api_key = cfg.get("api_key", "")
self._table_retrieval_dataset_id = (cfg.get("table_retrieval_dataset_id") or "").strip() 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() self._sql_gen_dataset_id = (cfg.get("sql_gen_dataset_id") or "").strip()
self._upload_path = (cfg.get("upload") or "").strip()
self._upload_mode = (cfg.get("upload_mode") or "overwrite").strip().lower()
def _validate_common(self) -> None: def _validate_common(self) -> None:
if not self._base_url: if not self._base_url:
raise RuntimeError("未配置 ragflow.url") raise RuntimeError("未配置 ragflow.url")
if not self._upload_path: if not self._upload_path:
raise RuntimeError("未配置 ragflow.upload 上传接口,请在 config/config.ini 中设置") 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"): if self._upload_mode not in ("overwrite", "append"):
raise RuntimeError("ragflow.upload_mode 仅支持 overwrite 或 append") raise RuntimeError("ragflow.upload_mode 仅支持 overwrite 或 append")
def _post(self, documents: List[Dict[str, any]]): def _post(self, documents: List[Dict[str, Any]], dataset_id: str):
self._validate_common() """上传文档到指定知识库 - 使用 multipart/form-data 格式"""
url = self._base_url + "/" + self._upload_path.lstrip("/") 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 {} headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {}
payload: Dict[str, any] = {"documents": documents}
if self._upload_mode == "overwrite": # 由于接口使用 multipart/form-data,我们需要创建临时文件
payload["mode"] = "overwrite" import tempfile
with httpx.Client(timeout=60) as client:
response = client.post(url, json=payload, headers=headers) # 创建临时文件并写入文档内容
response.raise_for_status() with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False, encoding='utf-8') as f:
return response.json() # 将文档内容写入文件
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 sync_table_retrieval(self) -> Dict[str, any]: 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")
# 构建正确的 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)
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}")
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")
url = f"{self._base_url}/api/v1/datasets/{dataset_id}"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {self._api_key}" if self._api_key else ""
}
print(f"更新知识库 URL: {url}") # 调试信息
print(f"更新配置: {config}") # 调试信息
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
def upload_table_retrieval(self) -> Dict[str, Any]:
"""上传表名检索模板文档 - 直接上传整个 JSON 文件"""
if not self._table_retrieval_dataset_id: if not self._table_retrieval_dataset_id:
raise RuntimeError("未配置 ragflow.table_retrieval_dataset_id,无法同步表名检索模板") raise RuntimeError("未配置 ragflow.table_retrieval_dataset_id,无法上传表名检索模板")
root = os.path.dirname(os.path.dirname(__file__)) root = os.path.dirname(os.path.dirname(__file__))
tables_file = os.path.join(root, "config", "table_retrieval_prompts", "tables.json") 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: with open(tables_file, "r", encoding="utf-8") as f:
data = json.load(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 = data.get("tables", {}) def upload_sql_gen(self) -> Dict[str, Any]:
documents = [] """上传 SQL 生成提示词文档"""
for table, templates in tables.items():
doc = {
"dataset_ids": self._table_retrieval_dataset_id,
"content": _build_document_for_table(table, templates),
"metadata": {"table": table},
}
documents.append(doc)
return self._post(documents)
def sync_sql_gen_prompts(self) -> Dict[str, any]:
if not self._sql_gen_dataset_id: if not self._sql_gen_dataset_id:
raise RuntimeError("未配置 ragflow.sql_gen_dataset_id,无法同步 SQL 生成提示词") raise RuntimeError("未配置 ragflow.sql_gen_dataset_id,无法上传 SQL 生成提示词")
root = os.path.dirname(os.path.dirname(__file__)) root = os.path.dirname(os.path.dirname(__file__))
prompts_dir = os.path.join(root, "config", "sql_gen_prompts") prompts_dir = os.path.join(root, "config", "sql_gen_prompts")
if not os.path.exists(prompts_dir):
raise RuntimeError(f"SQL 生成提示词目录不存在: {prompts_dir}")
documents = [] documents = []
for name in os.listdir(prompts_dir): for name in os.listdir(prompts_dir):
@@ -97,10 +248,12 @@ class RagflowSync:
prompt = json.load(f) prompt = json.load(f)
table = prompt.get("table") or os.path.splitext(name)[0] table = prompt.get("table") or os.path.splitext(name)[0]
doc = { doc = {
"dataset_ids": self._sql_gen_dataset_id,
"content": _build_sql_gen_document(table, prompt), "content": _build_sql_gen_document(table, prompt),
"metadata": {"table": table}, "metadata": {"table": table},
"title": f"SQL Prompt: {table}",
"type": "sql_prompt"
} }
documents.append(doc) documents.append(doc)
print(f"生成的 SQL 提示词文档: {doc}")
return self._post(documents) return self.upload_documents(self._sql_gen_dataset_id, documents)