276 lines
9.0 KiB
Python
276 lines
9.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Minimal Dify service API client for knowledge publication."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import dataclass
|
|
from typing import Final
|
|
from urllib.error import HTTPError, URLError
|
|
from urllib.parse import urlencode
|
|
from urllib.request import Request, urlopen
|
|
|
|
from .common import (
|
|
JSONMapping,
|
|
optional_non_empty_str,
|
|
require_mapping,
|
|
require_non_empty_str,
|
|
)
|
|
|
|
|
|
REQUEST_TIMEOUT_SECONDS: Final[int] = 15
|
|
|
|
|
|
class DifyClientError(Exception):
|
|
"""Raised when a Dify API request fails."""
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class DifyDataset:
|
|
"""Dify dataset summary."""
|
|
|
|
dataset_id: str
|
|
name: str
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class DifyDocument:
|
|
"""Dify document summary."""
|
|
|
|
document_id: str
|
|
name: str
|
|
|
|
|
|
class DifyClient:
|
|
"""Minimal service API client for Dify dataset operations."""
|
|
|
|
def __init__(self, base_url: str, api_key: str) -> None:
|
|
if not base_url.strip():
|
|
raise DifyClientError("缺少 Dify base URL")
|
|
if not api_key.strip():
|
|
raise DifyClientError("缺少 Dify API key")
|
|
self._base_url = base_url.rstrip("/")
|
|
self._api_key = api_key
|
|
|
|
def _request(
|
|
self,
|
|
method: str,
|
|
path: str,
|
|
*,
|
|
payload: JSONMapping | None = None,
|
|
query: dict[str, str] | None = None,
|
|
expect_empty: bool = False,
|
|
) -> JSONMapping:
|
|
url = f"{self._base_url}{path}"
|
|
if query:
|
|
url = f"{url}?{urlencode(query)}"
|
|
|
|
data: bytes | None = None
|
|
headers = {
|
|
"Authorization": f"Bearer {self._api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
if payload is not None:
|
|
data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
|
|
|
request = Request(url=url, data=data, method=method, headers=headers)
|
|
try:
|
|
with urlopen(request, timeout=REQUEST_TIMEOUT_SECONDS) as response:
|
|
raw_body = response.read().decode("utf-8").strip()
|
|
except HTTPError as exc:
|
|
error_body = exc.read().decode("utf-8", errors="replace")
|
|
raise DifyClientError(
|
|
f"Dify 请求失败: {method} {path} -> HTTP {exc.code}: {error_body}"
|
|
) from exc
|
|
except URLError as exc:
|
|
raise DifyClientError(
|
|
f"Dify 请求失败: {method} {path} -> {exc.reason}"
|
|
) from exc
|
|
|
|
if expect_empty and not raw_body:
|
|
return {}
|
|
if not raw_body:
|
|
return {}
|
|
try:
|
|
payload_object = json.loads(raw_body)
|
|
except json.JSONDecodeError as exc:
|
|
raise DifyClientError(f"Dify 返回了无效 JSON: {raw_body}") from exc
|
|
if expect_empty and not isinstance(payload_object, dict):
|
|
return {}
|
|
if not isinstance(payload_object, dict):
|
|
raise DifyClientError("Dify 返回根节点必须是 object")
|
|
return {str(key): value for key, value in payload_object.items()}
|
|
|
|
def list_datasets(self) -> list[DifyDataset]:
|
|
"""List Dify datasets visible to the service API token."""
|
|
|
|
payload = self._request("GET", "/datasets", query={"page": "1", "limit": "100"})
|
|
raw_items = payload.get("data", [])
|
|
if not isinstance(raw_items, list):
|
|
raise DifyClientError("/datasets.data 必须是列表")
|
|
|
|
results: list[DifyDataset] = []
|
|
for index, raw_item in enumerate(raw_items):
|
|
item = require_mapping(raw_item, f"/datasets.data[{index}]")
|
|
results.append(
|
|
DifyDataset(
|
|
dataset_id=require_non_empty_str(
|
|
item.get("id"), f"/datasets.data[{index}].id"
|
|
),
|
|
name=require_non_empty_str(
|
|
item.get("name"), f"/datasets.data[{index}].name"
|
|
),
|
|
)
|
|
)
|
|
return results
|
|
|
|
def ensure_dataset(
|
|
self, dataset_name: str, indexing_technique: str = "high_quality"
|
|
) -> DifyDataset:
|
|
"""Return an existing dataset or create it on demand."""
|
|
|
|
for dataset in self.list_datasets():
|
|
if dataset.name == dataset_name:
|
|
return dataset
|
|
|
|
payload = self._request(
|
|
"POST",
|
|
"/datasets",
|
|
payload={
|
|
"name": dataset_name,
|
|
"description": f"中央知识流水线管理的 {dataset_name} dataset",
|
|
"indexing_technique": indexing_technique,
|
|
"permission": "only_me",
|
|
},
|
|
)
|
|
dataset_id = require_non_empty_str(payload.get("id"), "POST /datasets.id")
|
|
created_name = require_non_empty_str(payload.get("name"), "POST /datasets.name")
|
|
return DifyDataset(dataset_id=dataset_id, name=created_name)
|
|
|
|
def list_documents(self, dataset_id: str) -> list[DifyDocument]:
|
|
"""List documents in one dataset."""
|
|
|
|
payload = self._request(
|
|
"GET",
|
|
f"/datasets/{dataset_id}/documents",
|
|
query={"page": "1", "limit": "100"},
|
|
)
|
|
raw_items = payload.get("data", [])
|
|
if not isinstance(raw_items, list):
|
|
raise DifyClientError("/documents.data 必须是列表")
|
|
|
|
results: list[DifyDocument] = []
|
|
for index, raw_item in enumerate(raw_items):
|
|
item = require_mapping(
|
|
raw_item, f"/datasets/{dataset_id}/documents.data[{index}]"
|
|
)
|
|
results.append(
|
|
DifyDocument(
|
|
document_id=require_non_empty_str(
|
|
item.get("id"),
|
|
f"/datasets/{dataset_id}/documents.data[{index}].id",
|
|
),
|
|
name=require_non_empty_str(
|
|
item.get("name"),
|
|
f"/datasets/{dataset_id}/documents.data[{index}].name",
|
|
),
|
|
)
|
|
)
|
|
return results
|
|
|
|
def find_document_id(self, dataset_id: str, document_name: str) -> str | None:
|
|
"""Find a document by its stable name."""
|
|
|
|
for document in self.list_documents(dataset_id):
|
|
if document.name == document_name:
|
|
return document.document_id
|
|
return None
|
|
|
|
def _create_document(
|
|
self,
|
|
*,
|
|
dataset_id: str,
|
|
document_name: str,
|
|
content: str,
|
|
indexing_technique: str,
|
|
) -> str:
|
|
"""Create one Dify text document."""
|
|
|
|
payload = self._request(
|
|
"POST",
|
|
f"/datasets/{dataset_id}/document/create_by_text",
|
|
payload={
|
|
"name": document_name,
|
|
"text": content,
|
|
"indexing_technique": indexing_technique,
|
|
"process_rule": {"mode": "automatic"},
|
|
},
|
|
)
|
|
raw_document = require_mapping(
|
|
payload.get("document"),
|
|
f"/datasets/{dataset_id}/document/create_by_text.document",
|
|
)
|
|
return require_non_empty_str(
|
|
raw_document.get("id"),
|
|
f"/datasets/{dataset_id}/document/create_by_text.document.id",
|
|
)
|
|
|
|
def upsert_document(
|
|
self,
|
|
*,
|
|
dataset_id: str,
|
|
document_name: str,
|
|
content: str,
|
|
document_id: str | None,
|
|
indexing_technique: str,
|
|
) -> str:
|
|
"""Create or update one Dify text document."""
|
|
|
|
if document_id:
|
|
try:
|
|
payload = self._request(
|
|
"POST",
|
|
f"/datasets/{dataset_id}/documents/{document_id}/update_by_text",
|
|
payload={
|
|
"name": document_name,
|
|
"text": content,
|
|
"process_rule": {"mode": "automatic"},
|
|
},
|
|
)
|
|
except DifyClientError as exc:
|
|
if "Document is not available" not in str(exc):
|
|
raise
|
|
return self._create_document(
|
|
dataset_id=dataset_id,
|
|
document_name=document_name,
|
|
content=content,
|
|
indexing_technique=indexing_technique,
|
|
)
|
|
updated_id = optional_non_empty_str(
|
|
payload.get("document_id"),
|
|
f"/datasets/{dataset_id}/documents/{document_id}/update_by_text.document_id",
|
|
)
|
|
return updated_id or document_id
|
|
|
|
return self._create_document(
|
|
dataset_id=dataset_id,
|
|
document_name=document_name,
|
|
content=content,
|
|
indexing_technique=indexing_technique,
|
|
)
|
|
|
|
def delete_document(self, dataset_id: str, document_id: str) -> None:
|
|
"""Delete one Dify document."""
|
|
|
|
try:
|
|
self._request(
|
|
"DELETE",
|
|
f"/datasets/{dataset_id}/documents/{document_id}",
|
|
expect_empty=True,
|
|
)
|
|
except DifyClientError as exc:
|
|
message = str(exc)
|
|
if "HTTP 404" in message or "Document Not Exists" in message:
|
|
return
|
|
raise
|