From a158cbbe9c4a94a6b6e299ed1c4ac6d0fdafb4ff Mon Sep 17 00:00:00 2001 From: chenjw28 <792430652@qq.com> Date: Thu, 26 Feb 2026 19:23:54 +0800 Subject: [PATCH] init --- api/endpoints.py | 47 +++- config/table_retrieval_prompts/tables.json | 45 ++-- services/ragflow_sync.py | 237 +++++++++++++++++---- 3 files changed, 260 insertions(+), 69 deletions(-) diff --git a/api/endpoints.py b/api/endpoints.py index a4d02b8..2945229 100644 --- a/api/endpoints.py +++ b/api/endpoints.py @@ -112,12 +112,49 @@ def reload_prompts(prompt_manager=Depends(get_prompt_manager)): @router.post("/api/ragflow/table-retrieval/reload") def reload_table_retrieval(): syncer = RagflowSync() - result = syncer.sync_table_retrieval() + result = syncer.upload_table_retrieval() return {"ok": True, "result": result} -@router.post("/api/ragflow/sql-gen/reload") -def reload_sql_gen(): +@router.post("/api/ragflow/table-retrieval/upload") +def upload_table_retrieval(): + """上传表名检索模板文档""" syncer = RagflowSync() - result = syncer.sync_sql_gen_prompts() - return {"ok": True, "result": result} + try: + 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)) diff --git a/config/table_retrieval_prompts/tables.json b/config/table_retrieval_prompts/tables.json index 65e0847..4223632 100644 --- a/config/table_retrieval_prompts/tables.json +++ b/config/table_retrieval_prompts/tables.json @@ -1,23 +1,24 @@ { - "tables": { - "apbo_eta_ful": [ - "SO ETA INFO", - "BO list", - "CC", - "Commodity Code", - "BO QTY", - "open order", - "BO status", - "BO report", - "model", - "MTM SN", - "backlog", - "Premier", - "SN machine_sn", - "work order information", - "recovery ETA", - "history order", - "Warranty type" - ] - } -} + "apbo_eta_ful": [ + "SO ETA INFO", + "BO list", + "CC", + "Commodity Code", + "BO QTY", + "open order", + "BO status", + "BO report", + "model", + "MTM SN", + "backlog", + "Premier", + "SN machine_sn", + "work order information", + "recovery ETA", + "history order", + "Warranty type" + ], + "x_table_name": [ + "xx" + ] +} \ No newline at end of file diff --git a/services/ragflow_sync.py b/services/ragflow_sync.py index 48dbd39..3a62c91 100644 --- a/services/ragflow_sync.py +++ b/services/ragflow_sync.py @@ -1,6 +1,6 @@ import json import os -from typing import Dict, List, Optional +from typing import Any, Dict, List import httpx @@ -8,25 +8,43 @@ from config import Config def _build_document_for_table(table: str, templates: List[str]) -> str: - lines = [f"table: {table}"] - for t in templates: - lines.append(f"- {t}") + """构建表名检索文档 - 使用更标准的格式""" + 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"table: {table}"] - if system_prompt: - lines.append("[system] " + system_prompt) - if business_prompt: - lines.append("[business] " + business_prompt) + + lines = [ + f"# SQL 生成提示词: {table}", + "", + "## 系统提示词:", + system_prompt, + "", + "## 业务提示词:", + business_prompt, + "" + ] + if constraints: - lines.append("[constraints]") - for c in constraints: - lines.append(f"- {c}") + 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) @@ -39,54 +57,187 @@ class RagflowSync: 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() - self._upload_path = (cfg.get("upload") or "").strip() - self._upload_mode = (cfg.get("upload_mode") or "overwrite").strip().lower() 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]]): - self._validate_common() - url = self._base_url + "/" + self._upload_path.lstrip("/") + 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 {} - payload: Dict[str, any] = {"documents": documents} - if self._upload_mode == "overwrite": - payload["mode"] = "overwrite" - with httpx.Client(timeout=60) as client: - response = client.post(url, json=payload, headers=headers) - response.raise_for_status() - return response.json() + + # 由于接口使用 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 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: - raise RuntimeError("未配置 ragflow.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]) - tables = data.get("tables", {}) - documents = [] - 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]: + def upload_sql_gen(self) -> Dict[str, Any]: + """上传 SQL 生成提示词文档""" 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__)) 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): @@ -97,10 +248,12 @@ class RagflowSync: prompt = json.load(f) table = prompt.get("table") or os.path.splitext(name)[0] doc = { - "dataset_ids": self._sql_gen_dataset_id, "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._post(documents) + return self.upload_documents(self._sql_gen_dataset_id, documents)