xiaozhi-esp32-server/main/manager-api-fastapi/app/integrations/ragflow.py

from __future__ import annotations

from collections.abc import Mapping
from typing import Any

import httpx
from fastapi import UploadFile

from app.core.errors import AppError

_DOCUMENT_CHUNK_METHODS = {
    "naive",
    "manual",
    "qa",
    "table",
    "paper",
    "book",
    "laws",
    "presentation",
    "picture",
    "one",
    "knowledge_graph",
    "email",
}
_RUN_STATUSES = {"UNSTART", "RUNNING", "CANCEL", "DONE", "FAIL"}
_DOCUMENT_PARSER_FIELDS = (
    "chunk_token_num",
    "delimiter",
    "layout_recognize",
    "html4excel",
    "auto_keywords",
    "auto_questions",
    "topn_tags",
    "raptor",
    "graphrag",
)
_DATASET_PARSER_FIELDS = (
    "chunk_token_num",
    "delimiter",
    "layout_recognize",
    "html4excel",
    "auto_keywords",
    "auto_questions",
)


class RAGFlowClient:
    """Async equivalent of the Java RAGFlow adapter and its wire contract."""

    def __init__(self, config: Mapping[str, Any]):
        self.config = dict(config)
        self.base_url = str(config.get("base_url") or config.get("baseUrl") or "").rstrip("/")
        self.api_key = str(config.get("api_key") or config.get("apiKey") or "")
        raw_timeout = config.get("timeout")
        if raw_timeout is None:
            self.timeout = 30.0
        else:
            try:
                self.timeout = float(int(str(raw_timeout)))
            except (TypeError, ValueError):
                self.timeout = 30.0
        self._validate(config)

    def _validate(self, config: Mapping[str, Any]) -> None:
        if not config:
            raise AppError(10164)
        if not self.base_url.strip():
            raise AppError(10171)
        if not self.api_key.strip():
            raise AppError(10172)
        if "你" in self.api_key:
            raise AppError(10173)
        if not self.base_url.startswith(("http://", "https://")):
            raise AppError(10174)
        adapter_type = "ragflow" if "type" not in config else str(config.get("type"))
        if adapter_type != "ragflow":
            raise AppError(10184, params=(f"适配器类型未注册: {adapter_type}",))

    async def request(
        self,
        method: str,
        endpoint: str,
        *,
        params: Mapping[str, Any] | None = None,
        json_body: Any = None,
        files: Mapping[str, Any] | None = None,
        data: Mapping[str, Any] | None = None,
    ) -> dict[str, Any]:
        headers = {"Authorization": f"Bearer {self.api_key}"}
        if files is None:
            headers["Content-Type"] = "application/json"
            headers["Accept-Charset"] = "utf-8"
        normalized_params = {
            key: self._query_value(value) for key, value in (params or {}).items() if value is not None
        }
        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                response = await client.request(
                    method,
                    self.base_url + endpoint,
                    params=normalized_params,
                    json=json_body,
                    files=files,
                    data=data,
                    headers=headers,
                )
                response.raise_for_status()
                payload = response.json()
        except (httpx.HTTPError, ValueError) as exc:
            raise AppError(10167, params=(f"Request Failed: {exc}",)) from exc
        if not isinstance(payload, dict):
            raise AppError(10167, params=("Invalid Response",))
        code = payload.get("code")
        if code is not None:
            if isinstance(code, bool) or not isinstance(code, int):
                raise AppError(10167, params=("Request Failed: invalid response code type",))
            if code != 0:
                message = payload.get("message")
                if message is not None and not isinstance(message, str):
                    raise AppError(10167, params=("Request Failed: invalid response message type",))
                raise AppError(10167, params=(message or "Unknown RAGFlow Error",))
        return dict(payload)

    @staticmethod
    def _query_value(value: Any) -> Any:
        if isinstance(value, bool):
            return str(value).lower()
        if isinstance(value, list):
            # Java List.toString() is what the baseline URL builder sends.
            return "[" + ", ".join(str(item) for item in value) + "]"
        return value

    async def dataset_info(self, dataset_id: str) -> dict[str, Any] | None:
        payload = await self.request(
            "GET", "/api/v1/datasets", params={"id": dataset_id, "page": 1, "page_size": 1}
        )
        data = payload.get("data")
        if isinstance(data, list) and data and isinstance(data[0], dict):
            return _normalize_dataset_info(data[0])
        return None

    async def create_dataset(self, body: dict[str, Any]) -> dict[str, Any]:
        body = dict(body)
        body["permission"] = "me" if _blank(body.get("permission")) else body.get("permission")
        body["chunk_method"] = "naive" if _blank(body.get("chunk_method")) else body.get("chunk_method")
        if _blank(body.get("embedding_model")):
            configured_model = self.config.get("embedding_model", self.config.get("embeddingModel"))
            body["embedding_model"] = None if _blank(configured_model) else configured_model
        body["avatar"] = body.get("avatar") if not _blank(body.get("avatar")) else (
            "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
        )
        body["parser_config"] = _normalize_parser_config(
            body.get("parser_config"), fields=_DATASET_PARSER_FIELDS
        )
        payload = await self.request("POST", "/api/v1/datasets", json_body=body)
        data = payload.get("data")
        if not isinstance(data, dict) or not data.get("id"):
            raise AppError(10167, params=("Invalid response from createDataset: missing data object",))
        return _normalize_dataset_info(data)

    async def update_dataset(self, dataset_id: str, body: dict[str, Any]) -> dict[str, Any] | None:
        body = dict(body)
        body["parser_config"] = _normalize_parser_config(
            body.get("parser_config"), fields=_DATASET_PARSER_FIELDS
        )
        payload = await self.request("PUT", f"/api/v1/datasets/{dataset_id}", json_body=body)
        return _normalize_dataset_info(payload["data"]) if isinstance(payload.get("data"), dict) else None

    async def delete_datasets(self, ids: list[str]) -> Any:
        return (await self.request("DELETE", "/api/v1/datasets", json_body={"ids": ids})).get("data")

    async def documents(
        self,
        dataset_id: str,
        *,
        page: int = 1,
        page_size: int = 10,
        name: str | None = None,
        status: str | None = None,
        document_id: str | None = None,
    ) -> tuple[list[dict[str, Any]], int]:
        params: dict[str, Any] = {"page": page, "page_size": page_size}
        if name:
            params["name"] = name
        if status:
            status_number = int(status) if status.lstrip("-").isdigit() else None
            names = {0: "UNSTART", 1: "RUNNING", 2: "CANCEL", 3: "DONE", 4: "FAIL"}
            params["run"] = [names[status_number]] if status_number in names else []
        if document_id:
            params["id"] = document_id
        payload = await self.request("GET", f"/api/v1/datasets/{dataset_id}/documents", params=params)
        data = payload.get("data")
        if not isinstance(data, dict):
            return [], 0
        docs = data.get("docs")
        rows: list[dict[str, Any]] = []
        if isinstance(docs, list):
            for item in docs:
                if not isinstance(item, dict):
                    continue
                try:
                    rows.append(_normalize_upload_document(item))
                except AppError:
                    # The Java adapter skips an individual document whose
                    # strong DTO conversion fails and keeps the rest of page.
                    continue
        return rows, int(data.get("total") or 0)

    async def upload_document(
        self,
        dataset_id: str,
        file: UploadFile,
        content: bytes,
        *,
        name: str,
        meta_fields: dict[str, Any] | None,
        chunk_method: str | None,
        parser_config: dict[str, Any] | None,
    ) -> dict[str, Any]:
        import json

        form: dict[str, Any] = {"name": name}
        if meta_fields:
            form["meta"] = json.dumps(meta_fields, ensure_ascii=False, separators=(",", ":"))
        if not _blank(chunk_method):
            normalized_method = str(chunk_method).lower()
            if normalized_method in _DOCUMENT_CHUNK_METHODS:
                form["chunk_method"] = normalized_method
        normalized_parser = (
            _normalize_parser_config(
                parser_config,
                fields=_DOCUMENT_PARSER_FIELDS,
                validate_layout=True,
            )
            if parser_config
            else None
        )
        if normalized_parser is not None:
            form["parser_config"] = json.dumps(normalized_parser, ensure_ascii=False, separators=(",", ":"))
        payload = await self.request(
            "POST",
            f"/api/v1/datasets/{dataset_id}/documents",
            files={"file": (file.filename or name, content, file.content_type or "application/octet-stream")},
            data=form,
        )
        data = payload.get("data")
        if isinstance(data, list) and data and isinstance(data[0], dict):
            return _normalize_upload_document(data[0])
        if isinstance(data, dict):
            return _normalize_upload_document(data)
        raise AppError(10167, params=("远程上传成功但未返回有效 DocumentID",))

    async def delete_documents(self, dataset_id: str, ids: list[str]) -> None:
        await self.request(
            "DELETE", f"/api/v1/datasets/{dataset_id}/documents", json_body={"ids": ids}
        )

    async def parse_documents(self, dataset_id: str, document_ids: list[str]) -> None:
        await self.request(
            "POST",
            f"/api/v1/datasets/{dataset_id}/chunks",
            json_body={"document_ids": document_ids},
        )

    async def chunks(
        self, dataset_id: str, document_id: str, params: Mapping[str, Any]
    ) -> dict[str, Any]:
        payload = await self.request(
            "GET", f"/api/v1/datasets/{dataset_id}/documents/{document_id}/chunks", params=params
        )
        data = payload.get("data")
        if not isinstance(data, dict):
            return {"chunks": [], "doc": None, "total": 0}
        try:
            return _normalize_chunk_list(data)
        except (TypeError, ValueError) as exc:
            raise AppError(10167, params=(str(exc),)) from exc

    async def retrieval(self, body: dict[str, Any]) -> dict[str, Any]:
        payload = await self.request("POST", "/api/v1/retrieval", json_body=body)
        data = payload.get("data")
        if not isinstance(data, dict):
            return {"chunks": [], "doc_aggs": [], "total": 0}
        try:
            return _normalize_retrieval_result(data)
        except (TypeError, ValueError) as exc:
            raise AppError(10167, params=(str(exc),)) from exc


def _blank(value: Any) -> bool:
    return value is None or (isinstance(value, str) and not value.strip())


def _normalize_parser_config(
    value: Any,
    *,
    fields: tuple[str, ...],
    validate_layout: bool = False,
) -> dict[str, Any] | None:
    if value is None:
        return None
    if not isinstance(value, Mapping):
        raise ValueError("parser_config must be an object")
    result = {key: value.get(key) for key in fields}
    if validate_layout and result.get("layout_recognize") not in {None, "DeepDOC", "Simple"}:
        raise ValueError("invalid layout_recognize")
    if "raptor" in result and result["raptor"] is not None:
        nested = result["raptor"]
        if not isinstance(nested, Mapping):
            raise ValueError("raptor must be an object")
        result["raptor"] = {"use_raptor": nested.get("use_raptor")}
    if "graphrag" in result and result["graphrag"] is not None:
        nested = result["graphrag"]
        if not isinstance(nested, Mapping):
            raise ValueError("graphrag must be an object")
        result["graphrag"] = {"use_graphrag": nested.get("use_graphrag")}
    return result


def _normalize_upload_document(value: Mapping[str, Any]) -> dict[str, Any]:
    result = dict(value)
    try:
        if result.get("parser_config") is not None:
            result["parser_config"] = _normalize_parser_config(
                result["parser_config"], fields=_DOCUMENT_PARSER_FIELDS, validate_layout=True
            )
    except (TypeError, ValueError) as exc:
        raise AppError(10167, params=("远程上传成功但未返回有效 DocumentID",)) from exc
    chunk_method = result.get("chunk_method")
    if chunk_method is not None:
        normalized_method = str(chunk_method).lower()
        if normalized_method not in _DOCUMENT_CHUNK_METHODS:
            raise AppError(10167, params=("远程上传成功但未返回有效 DocumentID",))
        result["chunk_method"] = normalized_method
    run = result.get("run")
    if run is not None and str(run) not in _RUN_STATUSES:
        raise AppError(10167, params=("远程上传成功但未返回有效 DocumentID",))
    return result


def _normalize_dataset_info(value: Mapping[str, Any]) -> dict[str, Any]:
    """Apply Jackson's DatasetDTO.InfoVO unknown-field and type boundary."""
    fields = (
        "id",
        "name",
        "avatar",
        "tenant_id",
        "description",
        "embedding_model",
        "permission",
        "chunk_method",
        "parser_config",
        "chunk_count",
        "document_count",
        "create_time",
        "update_time",
        "token_num",
        "create_date",
        "update_date",
    )
    try:
        result = {field: value.get(field) for field in fields}
        result["parser_config"] = _normalize_parser_config(
            result.get("parser_config"), fields=_DATASET_PARSER_FIELDS
        )
        for field in ("chunk_count", "document_count", "create_time", "update_time", "token_num"):
            raw_value = result[field]
            if raw_value is not None:
                result[field] = int(raw_value)
        return result
    except (TypeError, ValueError) as exc:
        raise AppError(10167, params=(str(exc),)) from exc


def _nullable_object(value: Any, fields: tuple[str, ...]) -> dict[str, Any] | None:
    if value is None:
        return None
    if not isinstance(value, Mapping):
        raise TypeError("response object has an invalid shape")
    return {field: value.get(field) for field in fields}


def _normalize_chunk_list(data: Mapping[str, Any]) -> dict[str, Any]:
    chunk_fields = (
        "id",
        "content",
        "document_id",
        "docnm_kwd",
        "important_keywords",
        "questions",
        "image_id",
        "dataset_id",
        "available",
        "positions",
        "token",
    )
    raw_chunks = data.get("chunks")
    if raw_chunks is None:
        chunks: list[dict[str, Any]] = []
    elif isinstance(raw_chunks, list):
        chunks = []
        for item in raw_chunks:
            normalized = _nullable_object(item, chunk_fields)
            if normalized is not None:
                chunks.append(normalized)
    else:
        raise TypeError("chunks must be an array")

    doc_fields = (
        "id",
        "thumbnail",
        "dataset_id",
        "chunk_method",
        "pipeline_id",
        "parser_config",
        "source_type",
        "type",
        "created_by",
        "name",
        "location",
        "size",
        "token_count",
        "chunk_count",
        "progress",
        "progress_msg",
        "process_begin_at",
        "process_duration",
        "meta_fields",
        "suffix",
        "run",
        "status",
        "create_time",
        "create_date",
        "update_time",
        "update_date",
    )
    doc = _nullable_object(data.get("doc"), doc_fields)
    if doc is not None:
        doc["parser_config"] = _normalize_parser_config(
            doc.get("parser_config"), fields=_DOCUMENT_PARSER_FIELDS, validate_layout=True
        )
        if doc.get("chunk_count") is not None:
            # DocumentDTO.InfoVO.chunkCount is Long, unlike the Integer field
            # on KnowledgeFilesDTO used by the document-list endpoint.
            doc["chunk_count"] = str(doc["chunk_count"])
        if doc.get("run") is not None and str(doc["run"]) not in _RUN_STATUSES:
            raise ValueError("invalid document run status")
    # ChunkDTO.ListVO.total is Long and therefore uses the Java global Long
    # serializer even for small values (including the adapter's default 0L).
    return {"chunks": chunks, "doc": doc, "total": str(int(data.get("total") or 0))}


def _normalize_retrieval_result(data: Mapping[str, Any]) -> dict[str, Any]:
    hit_fields = (
        "id",
        "content",
        "document_id",
        "dataset_id",
        "document_name",
        "document_keyword",
        "similarity",
        "vector_similarity",
        "term_similarity",
        "index",
        "highlight",
        "important_keywords",
        "questions",
        "image_id",
        "positions",
    )
    agg_fields = ("doc_name", "doc_id", "count")

    def normalize_list(raw: Any, fields: tuple[str, ...], name: str) -> list[dict[str, Any]]:
        if raw is None:
            return []
        if not isinstance(raw, list):
            raise TypeError(f"{name} must be an array")
        values: list[dict[str, Any]] = []
        for item in raw:
            normalized = _nullable_object(item, fields)
            if normalized is not None:
                values.append(normalized)
        return values

    return {
        "chunks": normalize_list(data.get("chunks"), hit_fields, "chunks"),
        "doc_aggs": normalize_list(data.get("doc_aggs"), agg_fields, "doc_aggs"),
        # RetrievalDTO.ResultVO.total is also Long.
        "total": str(int(data.get("total") or 0)),
    }