init

services/integrations/ragflow_client.py

@@ -0,0 +1,97 @@
+import json
+from typing import Any, Dict, Optional
+
+import httpx
+
+from config import Config
+
+
+TABLE_NAME_ALIASES = {
+    "apbo_tp_multiple_impact": "apbo_eta_multiple_impact",
+}
+
+
+class RagflowClient:
+    """RAGFlow 客户端（仅检索）"""
+
+    def __init__(self):
+        cfg = Config.get_section("ragflow")
+        self._base_url = cfg.get("url", "")
+        self._api_key = cfg.get("api_key", "")
+        self._retrieval_path = cfg.get("retrieval", "/api/v1/retrieval")
+        self._table_retrieval_dataset_id = cfg.get("table_retrieval_dataset_id", "")
+        self._sql_gen_dataset_id = cfg.get("sql_gen_dataset_id", "")
+
+    def _build_url(self) -> str:
+        return self._base_url.rstrip("/") + "/" + self._retrieval_path.lstrip("/")
+
+    @staticmethod
+    def _normalize_dataset_ids(dataset_id: Optional[str]) -> list[str]:
+        """将配置值规范化为 RAGFlow 需要的 list[string]"""
+        if not dataset_id:
+            return []
+        # 兼容逗号分隔配置
+        parts = [p.strip() for p in str(dataset_id).split(",") if p.strip()]
+        return parts
+
+    def retrieve(self, query: str, top_k: int = 3, dataset_id: Optional[str] = None, document_ids: Optional[str] = None) -> Dict[str, Any]:
+        """检索匹配文档"""
+        if not self._base_url or not self._retrieval_path:
+            raise RuntimeError("未配置 ragflow.url 或 ragflow.retrieval")
+        if not dataset_id:
+            raise ValueError("未提供 ragflow.dataset_id，无法进行检索")
+        url = self._build_url()
+        headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {}
+        payload = {
+            "dataset_ids": self._normalize_dataset_ids(dataset_id),
+            "question": query,
+            "top_k": top_k,
+        }
+        # 兼容部分版本字段
+        payload["query"] = query
+        if document_ids:
+            payload["document_ids"] = document_ids
+
+        with httpx.Client(timeout=30) as client:
+            response = client.post(url, json=payload, headers=headers)
+        response.raise_for_status()
+        return response.json()
+
+
+def extract_table_name(record: Dict[str, Any]) -> Optional[str]:
+    """从检索结果中提取表名"""
+    def _canonicalize(table_name: Any) -> Optional[str]:
+        normalized = str(table_name or "").strip()
+        if not normalized:
+            return None
+        return TABLE_NAME_ALIASES.get(normalized, normalized)
+
+    if not record:
+        return None
+
+    metadata = record.get("metadata") or {}
+    for key in ("table", "table_name"):
+        if key in metadata:
+            return _canonicalize(metadata.get(key))
+
+    for key in ("table", "table_name"):
+        if key in record:
+            return _canonicalize(record.get(key))
+
+    content = record.get("content") or record.get("text") or ""
+
+    # 兼容 content 为 JSON 字符串：{"table":"xxx", ...}
+    try:
+        parsed = json.loads(str(content))
+        if isinstance(parsed, dict):
+            for key in ("table", "table_name"):
+                if parsed.get(key):
+                    return _canonicalize(parsed.get(key))
+    except Exception:
+        pass
+
+    for line in str(content).splitlines():
+        if line.lower().startswith("table:"):
+            return _canonicalize(line.split(":", 1)[1].strip())
+
+    return None