import json from typing import Any, Dict, Optional import httpx from config import Config 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]: """从检索结果中提取表名""" if not record: return None metadata = record.get("metadata") or {} for key in ("table", "table_name"): if key in metadata: return metadata.get(key) for key in ("table", "table_name"): if key in record: return 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 str(parsed.get(key)) except Exception: pass for line in str(content).splitlines(): if line.lower().startswith("table:"): return line.split(":", 1)[1].strip() return None