def get_rerank_lists()

in hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py [0:0]

• 24 lines of code
• 5 McCabe index (conditional complexity)

    def get_rerank_lists(self, query: str, documents: List[str], top_n: Optional[int] = None) -> List[str]:
        if not top_n:
            top_n = len(documents)
        assert top_n <= len(documents), "'top_n' should be less than or equal to the number of documents"

        if top_n == 0:
            return []

        url = self.base_url
        from pyhugegraph.utils.constants import Constants
        headers = {
            "accept": Constants.HEADER_CONTENT_TYPE,
            "content-type": Constants.HEADER_CONTENT_TYPE,
            "Authorization": f"Bearer {self.api_key}",
        }
        payload = {
            "model": self.model,
            "query": query,
            "top_n": top_n,
            "documents": documents,
        }
        response = requests.post(url, headers=headers, json=payload, timeout=(1.0, 10.0))
        response.raise_for_status()  # Raise an error for bad status codes
        results = response.json()["results"]
        sorted_docs = [documents[item["index"]] for item in results]
        return sorted_docs