218 lines
8.9 KiB
Python
218 lines
8.9 KiB
Python
 
|
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 [])
|