ljwx-deploy/scripts/knowledge/dify_client.py

#!/usr/bin/env python3
"""Minimal Dify service API client for knowledge publication."""

from __future__ import annotations

import http.client
import json
import time
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] = 30

# Cloudflare in front of dify.omniverseai.net blocks the default
# `Python-urllib/*` User-Agent with error code 1010 (Browser Integrity
# Check). Set an explicit, identifiable UA to bypass the bot filter and
# leave a trace in Dify access logs for ops to attribute requests back to
# this client.
USER_AGENT: Final[str] = "ljwx-knowledge-client/1.0 (+https://github.com/lingjingwanxiang/ljwx-deploy)"

# Retry policy for transient infra hiccups: Dify gateway 5xx, mid-stream
# disconnect, generic connection refusal. Bulk ingests (hundreds of
# documents) saw single-request failures abort the whole run; without
# retry, operators had to manually replay the artifact. Exponential
# backoff cap chosen so a maxed-out retry sequence is still under a
# typical CI step's 5-minute idle timeout.
RETRY_MAX_ATTEMPTS: Final[int] = 4
RETRY_BACKOFF_SECONDS: Final[tuple[float, ...]] = (1.0, 3.0, 7.0)
RETRY_HTTP_STATUSES: Final[frozenset[int]] = frozenset({502, 503, 504})


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


@dataclass(frozen=True, slots=True)
class DifyRetrievalHit:
    """One scored result from a Dify retrieval call.

    `chunk_id` is the Dify-internal segment id; `document_id` is the Dify
    document UUID. Callers join these against publish-records to recover
    the original repo path and the stable `source_ref`.
    """

    chunk_id: str
    document_id: str
    document_name: str
    content: str
    score: float
    position: int


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")
        normalized_base_url = base_url.rstrip("/")
        if not normalized_base_url.endswith("/v1"):
            normalized_base_url = f"{normalized_base_url}/v1"
        self._base_url = normalized_base_url
        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",
            "User-Agent": USER_AGENT,
        }
        if payload is not None:
            data = json.dumps(payload, ensure_ascii=False).encode("utf-8")

        attempt = 0
        last_exc: Exception | None = None
        raw_body = ""
        while attempt < RETRY_MAX_ATTEMPTS:
            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()
                break
            except HTTPError as exc:
                error_body = exc.read().decode("utf-8", errors="replace")
                if exc.code in RETRY_HTTP_STATUSES and attempt + 1 < RETRY_MAX_ATTEMPTS:
                    last_exc = DifyClientError(
                        f"Dify 请求失败: {method} {path} -> HTTP {exc.code}: {error_body}"
                    )
                    delay = RETRY_BACKOFF_SECONDS[min(attempt, len(RETRY_BACKOFF_SECONDS) - 1)]
                    time.sleep(delay)
                    attempt += 1
                    continue
                raise DifyClientError(
                    f"Dify 请求失败: {method} {path} -> HTTP {exc.code}: {error_body}"
                ) from exc
            except (URLError, http.client.RemoteDisconnected, ConnectionError, TimeoutError) as exc:
                # urllib wraps RemoteDisconnected in URLError on some
                # Python versions and surfaces it bare on others; handle
                # both. Also retry on TimeoutError (raised when the
                # 30-second per-request budget is hit during Dify
                # backlog) and ConnectionError (TCP RST mid-stream).
                if attempt + 1 < RETRY_MAX_ATTEMPTS:
                    reason = getattr(exc, "reason", None) or exc.__class__.__name__
                    last_exc = DifyClientError(
                        f"Dify 请求失败: {method} {path} -> {reason}"
                    )
                    delay = RETRY_BACKOFF_SECONDS[min(attempt, len(RETRY_BACKOFF_SECONDS) - 1)]
                    time.sleep(delay)
                    attempt += 1
                    continue
                reason = getattr(exc, "reason", None) or exc.__class__.__name__
                raise DifyClientError(
                    f"Dify 请求失败: {method} {path} -> {reason}"
                ) from exc
        else:
            # Loop exited via `while` predicate becoming false without
            # break — every attempt raised something retryable. Surface
            # the last seen error so operators can see what Dify said.
            if last_exc is not None:
                raise last_exc
            raise DifyClientError(
                f"Dify 请求失败: {method} {path} -> exceeded {RETRY_MAX_ATTEMPTS} attempts"
            )

        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."""

        results: list[DifyDataset] = []
        page = 1
        while True:
            payload = self._request(
                "GET", "/datasets", query={"page": str(page), "limit": "100"}
            )
            raw_items = payload.get("data", [])
            if not isinstance(raw_items, list):
                raise DifyClientError("/datasets.data 必须是列表")

            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"
                        ),
                    )
                )
            if not payload.get("has_more") or not raw_items:
                break
            page += 1
        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."""

        results: list[DifyDocument] = []
        page = 1
        while True:
            payload = self._request(
                "GET",
                f"/datasets/{dataset_id}/documents",
                query={"page": str(page), "limit": "100"},
            )
            raw_items = payload.get("data", [])
            if not isinstance(raw_items, list):
                raise DifyClientError("/documents.data 必须是列表")

            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",
                        ),
                    )
                )
            if not payload.get("has_more") or not raw_items:
                break
            page += 1
        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:
                msg = str(exc)
                # Stale document_id (e.g. record migrated from old dify): fall
                # back to create. Catches both legacy "Document is not
                # available" wording and the HTTP 404 / "Document not found"
                # response from current Dify versions.
                is_stale = (
                    "Document is not available" in msg
                    or "Document not found" in msg
                    or "HTTP 404" in msg
                )
                if not is_stale:
                    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 retrieve(
        self,
        *,
        dataset_id: str,
        query: str,
        top_k: int = 10,
        score_threshold: float = 0.0,
        search_method: str = "semantic_search",
        reranking_enable: bool = False,
    ) -> list[DifyRetrievalHit]:
        """Retrieve chunks from one dataset for a query string.

        Maps to Dify Service API `POST /datasets/{id}/retrieve`. Returns
        the `records[]` flattened into typed hits, preserving the original
        position so callers can use it as a tiebreaker.
        """

        if not query.strip():
            raise DifyClientError("retrieve query 不能为空")
        if top_k < 1:
            raise DifyClientError("top_k 必须 >= 1")

        retrieval_model: JSONMapping = {
            "search_method": search_method,
            "reranking_enable": reranking_enable,
            "top_k": top_k,
            "score_threshold_enabled": score_threshold > 0,
            "score_threshold": score_threshold,
        }
        payload = self._request(
            "POST",
            f"/datasets/{dataset_id}/retrieve",
            payload={
                "query": query,
                "retrieval_model": retrieval_model,
            },
        )

        raw_records = payload.get("records", [])
        if not isinstance(raw_records, list):
            raise DifyClientError("/retrieve.records 必须是列表")

        hits: list[DifyRetrievalHit] = []
        for index, raw_record in enumerate(raw_records):
            record = require_mapping(raw_record, f"/retrieve.records[{index}]")
            segment = require_mapping(
                record.get("segment"), f"/retrieve.records[{index}].segment"
            )
            document = require_mapping(
                segment.get("document"),
                f"/retrieve.records[{index}].segment.document",
            )
            raw_score = record.get("score", 0.0)
            score = float(raw_score) if isinstance(raw_score, (int, float)) else 0.0
            hits.append(
                DifyRetrievalHit(
                    chunk_id=require_non_empty_str(
                        segment.get("id"),
                        f"/retrieve.records[{index}].segment.id",
                    ),
                    document_id=require_non_empty_str(
                        document.get("id"),
                        f"/retrieve.records[{index}].segment.document.id",
                    ),
                    document_name=require_non_empty_str(
                        document.get("name"),
                        f"/retrieve.records[{index}].segment.document.name",
                    ),
                    content=require_non_empty_str(
                        segment.get("content"),
                        f"/retrieve.records[{index}].segment.content",
                    ),
                    score=score,
                    position=index,
                )
            )
        return hits

    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