import json import os from typing import Any, Dict, List 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 _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) 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 _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") # 构建正确的 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: 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) # 将 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]) 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] 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}") return self.upload_documents(self._sql_gen_dataset_id, documents)