107 lines
4.1 KiB
Python
107 lines
4.1 KiB
Python
import json
|
|
import os
|
|
from typing import Dict, List, Optional
|
|
|
|
import httpx
|
|
|
|
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}")
|
|
return "\n".join(lines)
|
|
|
|
|
|
def _build_sql_gen_document(table: str, prompt: Dict[str, any]) -> str:
|
|
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)
|
|
if constraints:
|
|
lines.append("[constraints]")
|
|
for c in constraints:
|
|
lines.append(f"- {c}")
|
|
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()
|
|
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 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("/")
|
|
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()
|
|
|
|
def sync_table_retrieval(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 = 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]:
|
|
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 = []
|
|
|
|
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 = {
|
|
"dataset_ids": self._sql_gen_dataset_id,
|
|
"content": _build_sql_gen_document(table, prompt),
|
|
"metadata": {"table": table},
|
|
}
|
|
documents.append(doc)
|
|
|
|
return self._post(documents)
|