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 [])