fabricatio 0.2.1.dev0__cp313-cp313-win_amd64.whl → 0.3.14.dev4__cp313-cp313-win_amd64.whl

fabricatio/capabilities/rag.py

@@ -0,0 +1,263 @@
+"""A module for the RAG (Retrieval Augmented Generation) model."""
+
+from abc import ABC
+
+try:
+    from pymilvus import MilvusClient
+except ImportError as e:
+    raise RuntimeError(
+        "pymilvus is not installed. Have you installed `fabricatio[rag]` instead of `fabricatio`?"
+    ) from e
+from functools import lru_cache
+from operator import itemgetter
+from typing import List, Optional, Self, Type, Unpack
+
+from more_itertools.recipes import flatten, unique
+from pydantic import Field, PrivateAttr
+
+from fabricatio.journal import logger
+from fabricatio.models.adv_kwargs_types import CollectionConfigKwargs, FetchKwargs
+from fabricatio.models.extra.rag import MilvusDataBase
+from fabricatio.models.kwargs_types import ChooseKwargs
+from fabricatio.models.usages import EmbeddingUsage
+from fabricatio.rust import CONFIG, TEMPLATE_MANAGER
+from fabricatio.utils import ok
+
+
+@lru_cache(maxsize=None)
+def create_client(uri: str, token: str = "", timeout: Optional[float] = None) -> MilvusClient:
+    """Create a Milvus client."""
+    return MilvusClient(
+        uri=uri,
+        token=token,
+        timeout=timeout,
+    )
+
+
+class RAG(EmbeddingUsage, ABC):
+    """A class representing the RAG (Retrieval Augmented Generation) model."""
+
+    target_collection: Optional[str] = Field(default=None)
+    """The name of the collection being viewed."""
+
+    _client: Optional[MilvusClient] = PrivateAttr(None)
+    """The Milvus client used for the RAG model."""
+
+    @property
+    def client(self) -> MilvusClient:
+        """Return the Milvus client."""
+        if self._client is None:
+            raise RuntimeError("Client is not initialized. Have you called `self.init_client()`?")
+        return self._client
+
+    def init_client(
+        self,
+        milvus_uri: Optional[str] = None,
+        milvus_token: Optional[str] = None,
+        milvus_timeout: Optional[float] = None,
+    ) -> Self:
+        """Initialize the Milvus client."""
+        self._client = create_client(
+            uri=milvus_uri or ok(self.milvus_uri or CONFIG.rag.milvus_uri),
+            token=milvus_token
+            or (token.get_secret_value() if (token := (self.milvus_token or CONFIG.rag.milvus_token)) else ""),
+            timeout=milvus_timeout or self.milvus_timeout or CONFIG.rag.milvus_timeout,
+        )
+        return self
+
+    def check_client(self, init: bool = True) -> Self:
+        """Check if the client is initialized, and if not, initialize it."""
+        if self._client is None and init:
+            return self.init_client()
+        if self._client is None and not init:
+            raise RuntimeError("Client is not initialized. Have you called `self.init_client()`?")
+        return self
+
+    def view(
+        self, collection_name: Optional[str], create: bool = False, **kwargs: Unpack[CollectionConfigKwargs]
+    ) -> Self:
+        """View the specified collection.
+
+        Args:
+            collection_name (str): The name of the collection.
+            create (bool): Whether to create the collection if it does not exist.
+            **kwargs (Unpack[CollectionConfigKwargs]): Additional keyword arguments for collection configuration.
+        """
+        if create and collection_name and not self.check_client().client.has_collection(collection_name):
+            kwargs["dimension"] = ok(
+                kwargs.get("dimension")
+                or self.milvus_dimensions
+                or CONFIG.rag.milvus_dimensions
+                or self.embedding_dimensions
+                or CONFIG.embedding.dimensions,
+                "`dimension` is not set at any level.",
+            )
+            self.client.create_collection(collection_name, auto_id=True, **kwargs)
+            logger.info(f"Creating collection {collection_name}")
+
+        self.target_collection = collection_name
+        return self
+
+    def quit_viewing(self) -> Self:
+        """Quit the current view.
+
+        Returns:
+            Self: The current instance, allowing for method chaining.
+        """
+        return self.view(None)
+
+    @property
+    def safe_target_collection(self) -> str:
+        """Get the name of the collection being viewed, raise an error if not viewing any collection.
+
+        Returns:
+            str: The name of the collection being viewed.
+        """
+        return ok(self.target_collection, "No collection is being viewed. Have you called `self.view()`?")
+
+    async def add_document[D: MilvusDataBase](
+        self, data: List[D] | D, collection_name: Optional[str] = None, flush: bool = False
+    ) -> Self:
+        """Adds a document to the specified collection.
+
+        Args:
+            data (Union[Dict[str, Any], MilvusDataBase] | List[Union[Dict[str, Any], MilvusDataBase]]): The data to be added to the collection.
+            collection_name (Optional[str]): The name of the collection. If not provided, the currently viewed collection is used.
+            flush (bool): Whether to flush the collection after insertion.
+
+        Returns:
+            Self: The current instance, allowing for method chaining.
+        """
+        if isinstance(data, MilvusDataBase):
+            data = [data]
+
+        data_vec = await self.vectorize([d.prepare_vectorization() for d in data])
+        prepared_data = [d.prepare_insertion(vec) for d, vec in zip(data, data_vec, strict=True)]
+
+        c_name = collection_name or self.safe_target_collection
+        self.check_client().client.insert(c_name, prepared_data)
+
+        if flush:
+            logger.debug(f"Flushing collection {c_name}")
+            self.client.flush(c_name)
+        return self
+
+    async def afetch_document[D: MilvusDataBase](
+        self,
+        query: List[str],
+        document_model: Type[D],
+        collection_name: Optional[str] = None,
+        similarity_threshold: float = 0.37,
+        result_per_query: int = 10,
+        tei_endpoint: Optional[str] = None,
+        reranker_threshold: float = 0.7,
+        filter_expr: str = "",
+    ) -> List[D]:
+        """Asynchronously fetches documents from a Milvus database based on input vectors.
+
+        Args:
+           query (List[str]): A list of vectors to search for in the database.
+           document_model (Type[D]): The model class used to convert fetched data into document objects.
+           collection_name (Optional[str]): The name of the collection to search within.
+                                             If None, the currently viewed collection is used.
+           similarity_threshold (float): The similarity threshold for vector search. Defaults to 0.37.
+           result_per_query (int): The maximum number of results to return per query. Defaults to 10.
+           tei_endpoint (str): the endpoint of the TEI api.
+           reranker_threshold (float): The threshold used to filtered low relativity document.
+           filter_expr (str) : The filter expression used to filter out unwanted documents.
+
+        Returns:
+           List[D]: A list of document objects created from the fetched data.
+        """
+        # Step 1: Search for vectors
+        search_results = self.check_client().client.search(
+            collection_name or self.safe_target_collection,
+            await self.vectorize(query),
+            search_params={"radius": similarity_threshold},
+            output_fields=list(document_model.model_fields),
+            filter=filter_expr,
+            limit=result_per_query,
+        )
+        if tei_endpoint is not None:
+            from fabricatio.rust import TEIClient
+
+            reranker = TEIClient(base_url=tei_endpoint)
+
+            retrieved_id = set()
+            raw_result = []
+
+            for q, g in zip(query, search_results, strict=True):
+                models = document_model.from_sequence([res["entity"] for res in g if res["id"] not in retrieved_id])
+                logger.debug(f"Retrived {len(g)} raw document, filtered out {len(models)}.")
+                retrieved_id.update(res["id"] for res in g)
+                if not models:
+                    continue
+                rank_scores = await reranker.arerank(q, [m.prepare_vectorization() for m in models], truncate=True)
+                raw_result.extend((models[idx], scr) for (idx, scr) in rank_scores if scr > reranker_threshold)
+
+            raw_result_sorted = sorted(raw_result, key=lambda x: x[1], reverse=True)
+            return [r[0] for r in raw_result_sorted]
+
+        # Step 2: Flatten the search results
+        flattened_results = flatten(search_results)
+        unique_results = unique(flattened_results, key=itemgetter("id"))
+
+        # Step 3: Sort by distance (descending)
+        sorted_results = sorted(unique_results, key=itemgetter("distance"), reverse=True)
+
+        logger.debug(
+            f"Fetched {len(sorted_results)} document,searched similarities: {[t['distance'] for t in sorted_results]}"
+        )
+        # Step 4: Extract the entities
+        resp = [result["entity"] for result in sorted_results]
+
+        return document_model.from_sequence(resp)
+
+    async def aretrieve[D: MilvusDataBase](
+        self,
+        query: List[str] | str,
+        document_model: Type[D],
+        max_accepted: int = 20,
+        **kwargs: Unpack[FetchKwargs],
+    ) -> List[D]:
+        """Retrieve data from the collection.
+
+        Args:
+            query (List[str] | str): The query to be used for retrieval.
+            document_model (Type[D]): The model class used to convert retrieved data into document objects.
+            max_accepted (int): The final limit on the number of results to return.
+            **kwargs (Unpack[FetchKwargs]): Additional keyword arguments for retrieval.
+
+        Returns:
+            List[D]: A list of document objects created from the retrieved data.
+        """
+        if isinstance(query, str):
+            query = [query]
+
+        return (
+            await self.afetch_document(
+                query=query,
+                document_model=document_model,
+                **kwargs,
+            )
+        )[:max_accepted]
+
+    async def arefined_query(
+        self, question: List[str] | str, **kwargs: Unpack[ChooseKwargs[Optional[List[str]]]]
+    ) -> Optional[List[str]]:
+        """Refines the given question using a template.
+
+        Args:
+            question (List[str] | str): The question to be refined.
+            **kwargs (Unpack[ChooseKwargs]): Additional keyword arguments for the refinement process.
+
+        Returns:
+            List[str]: A list of refined questions.
+        """
+        return await self.alist_str(
+            TEMPLATE_MANAGER.render_template(
+                CONFIG.templates.refined_query_template,
+                {"question": [question] if isinstance(question, str) else question},
+            ),
+            **kwargs,
+        )

fabricatio/capabilities/rating.py

@@ -0,0 +1,404 @@
+"""A module that provides functionality to rate tasks based on a rating manual and score range."""
+
+from abc import ABC
+from itertools import permutations
+from random import sample
+from typing import Dict, List, Optional, Set, Tuple, Union, Unpack, overload
+
+from more_itertools import flatten, windowed
+from pydantic import Field, NonNegativeInt, PositiveInt, create_model
+
+from fabricatio.capabilities.propose import Propose
+from fabricatio.journal import logger
+from fabricatio.models.generic import Display, ProposedAble
+from fabricatio.models.kwargs_types import CompositeScoreKwargs, ValidateKwargs
+from fabricatio.parser import JsonCapture
+from fabricatio.rust import CONFIG, TEMPLATE_MANAGER
+from fabricatio.utils import ok, override_kwargs
+
+
+class Rating(Propose, ABC):
+    """A class that provides functionality to rate tasks based on a rating manual and score range.
+
+    References:
+        Lu X, Li J, Takeuchi K, et al. AHP-powered LLM reasoning for multi-criteria evaluation of open-ended responses[A/OL]. arXiv, 2024. DOI: 10.48550/arXiv.2410.01246.
+    """
+
+    async def rate_fine_grind(
+        self,
+        to_rate: str | List[str],
+        rating_manual: Dict[str, str],
+        score_range: Tuple[float, float],
+        **kwargs: Unpack[ValidateKwargs[Dict[str, float]]],
+    ) -> Dict[str, float] | List[Dict[str, float]] | List[Optional[Dict[str, float]]] | None:
+        """Rate a given string based on a rating manual and score range.
+
+        Args:
+            to_rate (str): The string to be rated.
+            rating_manual (Dict[str, str]): A dictionary containing the rating criteria.
+            score_range (Tuple[float, float]): A tuple representing the valid score range.
+            **kwargs (Unpack[ValidateKwargs]): Additional keyword arguments for the LLM usage.
+
+        Returns:
+            Dict[str, float]: A dictionary with the ratings for each dimension.
+        """
+        min_score, max_score = score_range
+        tip = (max_score - min_score) / 9
+
+        model = create_model(  # pyright: ignore [reportCallIssue]
+            "RatingResult",
+            __base__=ProposedAble,
+            __doc__=f"The rating result contains the scores against each criterion, with min_score={min_score} and max_score={max_score}.",
+            **{  # pyright: ignore [reportArgumentType]
+                criterion: (
+                    float,
+                    Field(
+                        ge=min_score,
+                        le=max_score,
+                        description=desc,
+                        examples=[round(min_score + tip * i, 2) for i in range(10)],
+                    ),
+                )
+                for criterion, desc in rating_manual.items()
+            },
+        )
+
+        res = await self.propose(
+            model,
+            TEMPLATE_MANAGER.render_template(
+                CONFIG.templates.rate_fine_grind_template,
+                {"to_rate": to_rate, "min_score": min_score, "max_score": max_score},
+            )
+            if isinstance(to_rate, str)
+            else [
+                TEMPLATE_MANAGER.render_template(
+                    CONFIG.templates.rate_fine_grind_template,
+                    {"to_rate": t, "min_score": min_score, "max_score": max_score},
+                )
+                for t in to_rate
+            ],
+            **override_kwargs(kwargs, default=None),
+        )
+        default = kwargs.get("default")
+        if isinstance(res, list):
+            return [r.model_dump() if r else default for r in res]
+        if res is None:
+            return default
+        return res.model_dump()
+
+    @overload
+    async def rate(
+        self,
+        to_rate: str,
+        topic: str,
+        criteria: Set[str],
+        manual: Optional[Dict[str, str]] = None,
+        score_range: Tuple[float, float] = (0.0, 1.0),
+        **kwargs: Unpack[ValidateKwargs],
+    ) -> Dict[str, float]: ...
+
+    @overload
+    async def rate(
+        self,
+        to_rate: List[str],
+        topic: str,
+        criteria: Set[str],
+        manual: Optional[Dict[str, str]] = None,
+        score_range: Tuple[float, float] = (0.0, 1.0),
+        **kwargs: Unpack[ValidateKwargs],
+    ) -> List[Dict[str, float]]: ...
+
+    async def rate(
+        self,
+        to_rate: Union[str, List[str]],
+        topic: str,
+        criteria: Set[str],
+        manual: Optional[Dict[str, str]] = None,
+        score_range: Tuple[float, float] = (0.0, 1.0),
+        **kwargs: Unpack[ValidateKwargs],
+    ) -> Dict[str, float] | List[Dict[str, float]] | List[Optional[Dict[str, float]]] | None:
+        """Rate a given string or a sequence of strings based on a topic, criteria, and score range.
+
+        Args:
+            to_rate (Union[str, List[str]]): The string or sequence of strings to be rated.
+            topic (str): The topic related to the task.
+            criteria (Set[str]): A set of criteria for rating.
+            manual (Optional[Dict[str, str]]): A dictionary containing the rating criteria. If not provided, then this method will draft the criteria automatically.
+            score_range (Tuple[float, float], optional): A tuple representing the valid score range. Defaults to (0.0, 1.0).
+            **kwargs (Unpack[ValidateKwargs]): Additional keyword arguments for the LLM usage.
+
+        Returns:
+            Union[Dict[str, float], List[Dict[str, float]]]: A dictionary with the ratings for each criterion if a single string is provided,
+            or a list of dictionaries with the ratings for each criterion if a sequence of strings is provided.
+        """
+        manual = (
+            manual
+            or await self.draft_rating_manual(topic, criteria, **override_kwargs(kwargs, default=None))
+            or dict(zip(criteria, criteria, strict=True))
+        )
+
+        return await self.rate_fine_grind(to_rate, manual, score_range, **kwargs)
+
+    async def draft_rating_manual(
+        self, topic: str, criteria: Optional[Set[str]] = None, **kwargs: Unpack[ValidateKwargs[Dict[str, str]]]
+    ) -> Optional[Dict[str, str]]:
+        """Drafts a rating manual based on a topic and dimensions.
+
+        Args:
+            topic (str): The topic for the rating manual.
+            criteria (Optional[Set[str]], optional): A set of criteria for the rating manual. If not specified, then this method will draft the criteria automatically.
+            **kwargs (Unpack[ValidateKwargs]): Additional keyword arguments for the LLM usage.
+
+        Returns:
+            Dict[str, str]: A dictionary representing the drafted rating manual.
+        """
+
+        def _validator(response: str) -> Dict[str, str] | None:
+            if (
+                (json_data := JsonCapture.validate_with(response, target_type=dict, elements_type=str)) is not None
+                and json_data.keys() == criteria
+                and all(isinstance(v, str) for v in json_data.values())
+            ):
+                return json_data
+            return None
+
+        criteria = criteria or await self.draft_rating_criteria(topic, **override_kwargs(dict(kwargs), default=None))
+
+        if criteria is None:
+            logger.error(f"Failed to draft rating criteria for topic {topic}")
+            return None
+
+        return await self.aask_validate(
+            question=(
+                TEMPLATE_MANAGER.render_template(
+                    CONFIG.templates.draft_rating_manual_template,
+                    {
+                        "topic": topic,
+                        "criteria": list(criteria),
+                    },
+                )
+            ),
+            validator=_validator,
+            **kwargs,
+        )
+
+    async def draft_rating_criteria(
+        self,
+        topic: str,
+        criteria_count: NonNegativeInt = 0,
+        **kwargs: Unpack[ValidateKwargs[Set[str]]],
+    ) -> Optional[Set[str]]:
+        """Drafts rating dimensions based on a topic.
+
+        Args:
+            topic (str): The topic for the rating dimensions.
+            criteria_count (NonNegativeInt, optional): The number of dimensions to draft, 0 means no limit. Defaults to 0.
+            **kwargs (Unpack[ValidateKwargs]): Additional keyword arguments for the LLM usage.
+
+        Returns:
+            Set[str]: A set of rating dimensions.
+        """
+        return await self.aask_validate(
+            question=(
+                TEMPLATE_MANAGER.render_template(
+                    CONFIG.templates.draft_rating_criteria_template,
+                    {
+                        "topic": topic,
+                        "criteria_count": criteria_count,
+                    },
+                )
+            ),
+            validator=lambda resp: set(out)
+            if (out := JsonCapture.validate_with(resp, list, str, criteria_count)) is not None
+            else out,
+            **kwargs,
+        )
+
+    async def draft_rating_criteria_from_examples(
+        self,
+        topic: str,
+        examples: List[str],
+        m: NonNegativeInt = 0,
+        reasons_count: PositiveInt = 2,
+        criteria_count: PositiveInt = 5,
+        **kwargs: Unpack[ValidateKwargs],
+    ) -> Optional[Set[str]]:
+        """Asynchronously drafts a set of rating criteria based on provided examples.
+
+        This function generates rating criteria by analyzing examples and extracting reasons for comparison,
+        then further condensing these reasons into a specified number of criteria.
+
+        Parameters:
+            topic (str): The subject topic for the rating criteria.
+            examples (List[str]): A list of example texts to analyze.
+            m (NonNegativeInt, optional): The number of examples to sample from the provided list. Defaults to 0 (no sampling).
+            reasons_count (PositiveInt, optional): The number of reasons to extract from each pair of examples. Defaults to 2.
+            criteria_count (PositiveInt, optional): The final number of rating criteria to draft. Defaults to 5.
+            **kwargs (Unpack[ValidateKwargs]): Additional keyword arguments for validation.
+
+        Returns:
+            Set[str]: A set of drafted rating criteria.
+
+        Warnings:
+            Since this function uses pairwise comparisons, it may not be suitable for large lists of examples.
+            For that reason, consider using a smaller list of examples or setting `m` to a non-zero value smaller than the length of the examples.
+        """
+        if m:
+            examples = sample(examples, m)
+
+        # extract reasons from the comparison of ordered pairs of extracted from examples
+        reasons = flatten(
+            await self.aask_validate(  # pyright: ignore [reportArgumentType]
+                question=[
+                    TEMPLATE_MANAGER.render_template(
+                        CONFIG.templates.extract_reasons_from_examples_template,
+                        {
+                            "topic": topic,
+                            "first": pair[0],
+                            "second": pair[1],
+                            "reasons_count": reasons_count,
+                        },
+                    )
+                    for pair in (permutations(examples, 2))
+                ],
+                validator=lambda resp: JsonCapture.validate_with(
+                    resp, target_type=list, elements_type=str, length=reasons_count
+                ),
+                **kwargs,
+            )
+        )
+        # extract certain mount of criteria from reasons according to their importance and frequency
+        return await self.aask_validate(
+            question=(
+                TEMPLATE_MANAGER.render_template(
+                    CONFIG.templates.extract_criteria_from_reasons_template,
+                    {
+                        "topic": topic,
+                        "reasons": list(reasons),
+                        "criteria_count": criteria_count,
+                    },
+                )
+            ),
+            validator=lambda resp: set(out)
+            if (out := JsonCapture.validate_with(resp, target_type=list, elements_type=str, length=criteria_count))
+            else None,
+            **kwargs,
+        )
+
+    async def drafting_rating_weights_klee(
+        self,
+        topic: str,
+        criteria: Set[str],
+        **kwargs: Unpack[ValidateKwargs[float]],
+    ) -> Dict[str, float]:
+        """Drafts rating weights for a given topic and criteria using the Klee method.
+
+        Args:
+            topic (str): The topic for the rating weights.
+            criteria (Set[str]): A set of criteria for the rating weights.
+            **kwargs (Unpack[ValidateKwargs]): Additional keyword arguments for the LLM usage.
+
+        Returns:
+            Dict[str, float]: A dictionary representing the drafted rating weights for each criterion.
+        """
+        if len(criteria) < 2:  # noqa: PLR2004
+            raise ValueError("At least two criteria are required to draft rating weights")
+
+        criteria_seq = list(criteria)  # freeze the order
+        windows = windowed(criteria_seq, 2)
+
+        # get the importance multiplier indicating how important is second criterion compared to the first one
+        relative_weights = await self.aask_validate(
+            question=[
+                TEMPLATE_MANAGER.render_template(
+                    CONFIG.templates.draft_rating_weights_klee_template,
+                    {
+                        "topic": topic,
+                        "first": pair[0],
+                        "second": pair[1],
+                    },
+                )
+                for pair in windows
+            ],
+            validator=lambda resp: JsonCapture.validate_with(resp, target_type=float),
+            **kwargs,
+        )
+        if not all(relative_weights):
+            raise ValueError(f"found illegal weight: {relative_weights}")
+        weights = [1.0]
+        for rw in relative_weights:
+            weights.append(weights[-1] * rw)  # pyright: ignore [reportOperatorIssue]
+        total = sum(weights)
+        return dict(zip(criteria_seq, [w / total for w in weights], strict=True))
+
+    async def composite_score(
+        self,
+        topic: str,
+        to_rate: List[str],
+        criteria: Optional[Set[str]] = None,
+        weights: Optional[Dict[str, float]] = None,
+        manual: Optional[Dict[str, str]] = None,
+        approx: bool = False,
+        **kwargs: Unpack[ValidateKwargs[List[Dict[str, float]]]],
+    ) -> List[float]:
+        """Calculates the composite scores for a list of items based on a given topic and criteria.
+
+        Args:
+            topic (str): The topic for the rating.
+            to_rate (List[str]): A list of strings to be rated.
+            criteria (Optional[Set[str]]): A set of criteria for the rating. Defaults to None.
+            weights (Optional[Dict[str, float]]): A dictionary of rating weights for each criterion. Defaults to None.
+            manual (Optional[Dict[str, str]]): A dictionary of manual ratings for each item. Defaults to None.
+            approx (bool): Whether to use approximate rating criteria. Defaults to False.
+            **kwargs (Unpack[ValidateKwargs]): Additional keyword arguments for the LLM usage.
+
+        Returns:
+            List[float]: A list of composite scores for the items.
+        """
+        criteria = ok(
+            criteria
+            or (await self.draft_rating_criteria(topic, **override_kwargs(kwargs, default=None)) if approx else None)
+            or await self.draft_rating_criteria_from_examples(topic, to_rate, **override_kwargs(kwargs, default=None))
+        )
+        weights = ok(
+            weights or await self.drafting_rating_weights_klee(topic, criteria, **override_kwargs(kwargs, default=None))
+        )
+        logger.info(f"Criteria: {criteria}\nWeights: {weights}")
+        ratings_seq = await self.rate(to_rate, topic, criteria, manual, **kwargs)
+
+        return [sum(ratings[c] * weights[c] for c in criteria) for ratings in ratings_seq]
+
+    @overload
+    async def best(self, candidates: List[str], k: int = 1, **kwargs: Unpack[CompositeScoreKwargs]) -> List[str]: ...
+
+    @overload
+    async def best[T: Display](
+        self, candidates: List[T], k: int = 1, **kwargs: Unpack[CompositeScoreKwargs]
+    ) -> List[T]: ...
+
+    async def best[T: Display](
+        self, candidates: List[str] | List[T], k: int = 1, **kwargs: Unpack[CompositeScoreKwargs]
+    ) -> Optional[List[str] | List[T]]:
+        """Choose the best candidates from the list of candidates based on the composite score.
+
+        Args:
+            k (int): The number of best candidates to choose.
+            candidates (List[str]): A list of candidates to choose from.
+            **kwargs (CompositeScoreKwargs): Additional keyword arguments for the composite score calculation.
+
+        Returns:
+            List[str]: The best candidates.
+        """
+        if (leng := len(candidates)) == 0:
+            logger.warning(f"No candidates, got {leng}, return None.")
+            return None
+
+        if leng == 1:
+            logger.warning(f"Only one candidate, got {leng}, return it.")
+            return candidates
+        logger.info(f"Choose best {k} from {leng} candidates.")
+
+        rating_seq = await self.composite_score(
+            to_rate=[c.display() if isinstance(c, Display) else c for c in candidates], **kwargs
+        )
+        return [a[0] for a in sorted(zip(candidates, rating_seq, strict=True), key=lambda x: x[1], reverse=True)[:k]]  # pyright: ignore [reportReturnType]