98 lines
3.3 KiB
Python
98 lines
3.3 KiB
Python
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
|