dify/api/controllers/console/datasets/hit_testing_base.py

import logging
from typing import Any, cast

from pydantic import BaseModel, Field
from sqlalchemy.orm import Session
from werkzeug.exceptions import Forbidden, InternalServerError, NotFound

import services.errors.base
from controllers.console.app.error import (
    CompletionRequestError,
    ProviderModelCurrentlyNotSupportError,
    ProviderNotInitializeError,
    ProviderQuotaExceededError,
)
from controllers.console.datasets.error import DatasetNotInitializedError
from core.errors.error import (
    LLMBadRequestError,
    ModelCurrentlyNotSupportError,
    ProviderTokenNotInitError,
    QuotaExceededError,
)
from graphon.model_runtime.errors.invoke import InvokeError
from libs.login import resolve_account_fallback
from models.account import Account
from models.dataset import Dataset
from services.hit_testing_service import HitTestingService
from services.knowledge.dataset_service import DatasetService
from services.knowledge.entities.knowledge_entities import ExternalRetrievalModel, RetrievalModel

logger = logging.getLogger(__name__)


class HitTestingPayload(BaseModel):
    query: str = Field(description="Search query text.", max_length=250)
    retrieval_model: RetrievalModel | None = Field(
        default=None,
        description=(
            "Retrieval model configuration. Controls how chunks are searched and ranked when querying this "
            "knowledge base."
        ),
    )
    external_retrieval_model: ExternalRetrievalModel = Field(
        default=None,
        description="Retrieval settings for external knowledge bases.",
    )
    attachment_ids: list[str] | None = Field(
        default=None,
        description="List of attachment IDs to include in the retrieval context.",
    )


class DatasetsHitTestingBase:
    @staticmethod
    def _prepare_hit_testing_records(records: Any) -> list[dict[str, Any]]:
        """Ensure collection fields match the API schema before response validation."""
        if not isinstance(records, list):
            raise ValueError("Invalid hit testing records response")

        normalized_records: list[dict[str, Any]] = []
        for record in records:
            if not isinstance(record, dict):
                raise ValueError("Invalid hit testing record response")

            normalized_record = dict(record)
            segment = normalized_record.get("segment")
            if isinstance(segment, dict):
                normalized_segment = dict(segment)
                normalized_segment.setdefault("sign_content", None)
                if normalized_segment.get("keywords") is None:
                    normalized_segment["keywords"] = []
                normalized_record["segment"] = normalized_segment

            if normalized_record.get("child_chunks") is None:
                normalized_record["child_chunks"] = []

            if normalized_record.get("files") is None:
                normalized_record["files"] = []

            normalized_record.setdefault("tsne_position", None)
            normalized_record.setdefault("summary", None)

            normalized_records.append(normalized_record)

        return normalized_records

    @staticmethod
    def get_and_validate_dataset(
        session: Session,
        dataset_id: str,
        current_user: Account | None = None,
        current_tenant_id: str | None = None,
    ) -> Dataset:
        current_user, _ = resolve_account_fallback(current_user, current_tenant_id)
        dataset = DatasetService.get_dataset(dataset_id, session)
        if dataset is None:
            raise NotFound("Dataset not found.")

        try:
            DatasetService.check_dataset_permission(dataset, current_user, session)
        except services.errors.base.NoPermissionError as e:
            raise Forbidden(str(e)) from e

        return dataset

    @staticmethod
    def hit_testing_args_check(args: dict[str, Any]) -> None:
        HitTestingService.hit_testing_args_check(args)

    @staticmethod
    def parse_args(payload: dict[str, Any] | None) -> dict[str, Any]:
        """Validate and return hit-testing arguments from an incoming payload."""
        hit_testing_payload = HitTestingPayload.model_validate(payload or {})
        return hit_testing_payload.model_dump(exclude_none=True)

    @staticmethod
    def perform_hit_testing(
        session: Session,
        dataset: Dataset,
        args: dict[str, Any],
        current_user: Account | None = None,
        current_tenant_id: str | None = None,
    ) -> dict[str, Any]:
        try:
            current_user, _ = resolve_account_fallback(current_user, current_tenant_id)
            response = HitTestingService.retrieve(
                session=session,
                dataset=dataset,
                query=cast(str, args.get("query")),
                account=current_user,
                retrieval_model=args.get("retrieval_model"),
                external_retrieval_model=cast(dict[str, Any], args.get("external_retrieval_model")),
                attachment_ids=args.get("attachment_ids"),
                limit=10,
            )
            query = response.get("query")
            if not isinstance(query, dict) or not isinstance(query.get("content"), str):
                raise ValueError("Invalid hit testing query response")

            return {
                "query": {"content": query["content"]},
                "records": DatasetsHitTestingBase._prepare_hit_testing_records(response.get("records", [])),
            }
        except services.errors.index.IndexNotInitializedError as caught_error:
            raise DatasetNotInitializedError() from caught_error
        except ProviderTokenNotInitError as ex:
            raise ProviderNotInitializeError(ex.description) from ex
        except QuotaExceededError as caught_error:
            raise ProviderQuotaExceededError() from caught_error
        except ModelCurrentlyNotSupportError as caught_error:
            raise ProviderModelCurrentlyNotSupportError() from caught_error
        except LLMBadRequestError as caught_error:
            raise ProviderNotInitializeError(
                "No Embedding Model or Reranking Model available. Please configure a valid provider "
                "in the Settings -> Model Provider."
            ) from caught_error
        except InvokeError as e:
            raise CompletionRequestError(e.description) from e
        except ValueError as e:
            raise ValueError(str(e)) from e
        except Exception as e:
            logger.exception("Hit testing failed.")
            raise InternalServerError() from e