feat: 用 Ragflow 搭建鲜境方案客服

咨询台接 Ragflow 检索与示例知识库，支持追问、判定、引用和留资。 Co-authored-by: haijd <haijd@163.com>
2026-09-20 01:08:11 +00:00
parent 9c6afcb61a
commit b7bd830da7
20 changed files with 925 additions and 258 deletions
@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+import base64
+from pathlib import Path
+from typing import Any
+
+import httpx
+try:
+    from Cryptodome.Cipher import PKCS1_v1_5
+    from Cryptodome.PublicKey import RSA
+except ImportError:
+    from Crypto.Cipher import PKCS1_v1_5
+    from Crypto.PublicKey import RSA
+
+from .settings import ROOT, settings
+
+PUBLIC_PEM = ROOT / "infra" / "ragflow-public.pem"
+
+
+class RagflowError(RuntimeError):
+    pass
+
+
+def encrypt_password(password: str) -> str:
+    """Match Ragflow crypt(): RSA-encrypt base64(password)."""
+    rsa_key = RSA.import_key(PUBLIC_PEM.read_text())
+    cipher = PKCS1_v1_5.new(rsa_key)
+    password_b64 = base64.b64encode(password.encode("utf-8"))
+    encrypted = cipher.encrypt(password_b64)
+    return base64.b64encode(encrypted).decode("utf-8")
+
+
+class RagflowClient:
+    def __init__(self, base_url: str | None = None, api_key: str | None = None) -> None:
+        self.base_url = (base_url or settings.ragflow_base_url).rstrip("/")
+        self.api_key = api_key if api_key is not None else settings.ragflow_api_key
+
+    def _headers(self, token: str | None = None) -> dict[str, str]:
+        key = token or self.api_key
+        if not key:
+            return {}
+        if key.lower().startswith("bearer "):
+            return {"Authorization": key}
+        return {"Authorization": f"Bearer {key}"}
+
+    def health(self) -> dict[str, Any]:
+        url = f"{self.base_url}/api/v1/system/healthz"
+        try:
+            with httpx.Client(timeout=8.0) as client:
+                resp = client.get(url)
+            ok = resp.status_code < 500
+            return {"ok": ok, "status_code": resp.status_code, "url": url}
+        except Exception as exc:
+            return {"ok": False, "error": str(exc), "url": url}
+
+    def register(self, email: str, password: str, nickname: str = "鲜境管理员") -> str:
+        payload = {
+            "nickname": nickname,
+            "email": email,
+            "password": encrypt_password(password),
+        }
+        with httpx.Client(timeout=30.0) as client:
+            resp = client.post(f"{self.base_url}/api/v1/users", json=payload)
+        if resp.status_code >= 400:
+            raise RagflowError(f"注册失败 {resp.status_code}: {resp.text[:400]}")
+        auth = resp.headers.get("Authorization") or ""
+        body = resp.json() if resp.content else {}
+        if body.get("code") not in (0, None) and body.get("code") != 0:
+            # already registered or validation error
+            message = str(body.get("message") or body)
+            if "already" in message.lower() or "已经" in message:
+                return self.login(email, password)
+            raise RagflowError(message)
+        if not auth:
+            data = body.get("data") or {}
+            auth = str(data.get("id") or "")
+        if not auth:
+            raise RagflowError(f"注册成功但没有登录凭证: {body}")
+        return auth
+
+    def login(self, email: str, password: str) -> str:
+        payload = {"email": email, "password": encrypt_password(password)}
+        with httpx.Client(timeout=30.0) as client:
+            resp = client.post(f"{self.base_url}/api/v1/auth/login", json=payload)
+        if resp.status_code >= 400:
+            raise RagflowError(f"登录失败 {resp.status_code}: {resp.text[:400]}")
+        auth = resp.headers.get("Authorization") or ""
+        if not auth:
+            raise RagflowError("登录没有返回 Authorization")
+        return auth
+
+    def ensure_api_token(self, session_auth: str) -> str:
+        headers = {"Authorization": session_auth}
+        with httpx.Client(timeout=30.0) as client:
+            listed = client.get(f"{self.base_url}/api/v1/system/tokens", headers=headers)
+            if listed.status_code < 400:
+                data = (listed.json() or {}).get("data") or []
+                if isinstance(data, list) and data:
+                    token = data[0].get("token")
+                    if token:
+                        return token
+            created = client.post(f"{self.base_url}/api/v1/system/tokens", headers=headers)
+        if created.status_code >= 400:
+            raise RagflowError(f"创建 API Token 失败: {created.text[:400]}")
+        body = created.json()
+        token = (body.get("data") or {}).get("token")
+        if not token:
+            raise RagflowError(f"Token 响应异常: {body}")
+        return token
+
+    def list_datasets(self) -> list[dict[str, Any]]:
+        with httpx.Client(timeout=20.0) as client:
+            resp = client.get(f"{self.base_url}/api/v1/datasets", headers=self._headers(), params={"page": 1, "page_size": 50})
+        resp.raise_for_status()
+        data = (resp.json() or {}).get("data") or []
+        if isinstance(data, dict):
+            data = data.get("datasets") or data.get("kbs") or []
+        return list(data)
+
+    def create_dataset(self, name: str, description: str) -> dict[str, Any]:
+        payload = {
+            "name": name,
+            "description": description,
+            "chunk_method": "naive",
+            "permission": "me",
+            "embedding_model": "BAAI/bge-small-zh-v1.5@Builtin",
+            "parser_config": {
+                "graphrag": {"use_graphrag": False},
+                "raptor": {"use_raptor": False},
+            },
+        }
+        with httpx.Client(timeout=30.0) as client:
+            resp = client.post(f"{self.base_url}/api/v1/datasets", headers=self._headers(), json=payload)
+        if resp.status_code >= 400:
+            raise RagflowError(resp.text[:400])
+        body = resp.json()
+        if body.get("code") not in (0, None):
+            raise RagflowError(str(body.get("message") or body))
+        return body.get("data") or {}
+
+    def update_dataset(self, dataset_id: str, payload: dict[str, Any]) -> dict[str, Any]:
+        with httpx.Client(timeout=30.0) as client:
+            resp = client.put(
+                f"{self.base_url}/api/v1/datasets/{dataset_id}",
+                headers=self._headers(),
+                json=payload,
+            )
+        if resp.status_code >= 400:
+            raise RagflowError(resp.text[:400])
+        body = resp.json()
+        if body.get("code") not in (0, None):
+            raise RagflowError(str(body.get("message") or body))
+        return body.get("data") or {}
+
+    def upload_document(self, dataset_id: str, path: Path) -> list[dict[str, Any]]:
+        with path.open("rb") as fh:
+            files = {"file": (path.name, fh, "text/markdown")}
+            with httpx.Client(timeout=60.0) as client:
+                resp = client.post(
+                    f"{self.base_url}/api/v1/datasets/{dataset_id}/documents",
+                    headers=self._headers(),
+                    files=files,
+                )
+        if resp.status_code >= 400:
+            raise RagflowError(resp.text[:500])
+        body = resp.json()
+        data = body.get("data") or []
+        if isinstance(data, dict):
+            data = [data]
+        return data
+
+    def parse_documents(self, dataset_id: str, document_ids: list[str]) -> None:
+        with httpx.Client(timeout=30.0) as client:
+            resp = client.post(
+                f"{self.base_url}/api/v1/datasets/{dataset_id}/documents/parse",
+                headers=self._headers(),
+                json={"document_ids": document_ids},
+            )
+        if resp.status_code >= 400:
+            raise RagflowError(resp.text[:400])
+
+    def list_documents(self, dataset_id: str) -> list[dict[str, Any]]:
+        with httpx.Client(timeout=20.0) as client:
+            resp = client.get(
+                f"{self.base_url}/api/v1/datasets/{dataset_id}/documents",
+                headers=self._headers(),
+                params={"page": 1, "page_size": 50},
+            )
+        resp.raise_for_status()
+        data = (resp.json() or {}).get("data") or []
+        if isinstance(data, dict):
+            data = data.get("docs") or data.get("documents") or []
+        return list(data)
+
+    def retrieve(self, question: str, dataset_ids: list[str], page_size: int = 4) -> list[dict[str, Any]]:
+        if not self.api_key:
+            raise RagflowError("未配置 RAGFLOW_API_KEY")
+        if not dataset_ids:
+            raise RagflowError("未配置 RAGFLOW_DATASET_IDS")
+        payload = {
+            "question": question,
+            "dataset_ids": dataset_ids,
+            "page": 1,
+            "page_size": page_size,
+            "similarity_threshold": 0.15,
+            "vector_similarity_weight": 0.3,
+        }
+        with httpx.Client(timeout=30.0) as client:
+            resp = client.post(f"{self.base_url}/api/v1/retrieval", headers=self._headers(), json=payload)
+        if resp.status_code >= 400:
+            raise RagflowError(f"检索失败 {resp.status_code}: {resp.text[:300]}")
+        body = resp.json()
+        if body.get("code") not in (0, None):
+            raise RagflowError(str(body.get("message") or body))
+        data = body.get("data") or {}
+        chunks = data.get("chunks") if isinstance(data, dict) else data
+        return list(chunks or [])