def get_rerank_lists()

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

• 27 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 = "https://api.siliconflow.cn/v1/rerank"
        payload = {
            "model": self.model,
            "query": query,
            "documents": documents,
            "return_documents": False,
            "max_chunks_per_doc": 1024,
            "overlap_tokens": 80,
            "top_n": top_n,
        }
        from pyhugegraph.utils.constants import Constants
        headers = {
            "accept": Constants.HEADER_CONTENT_TYPE,
            "content-type": Constants.HEADER_CONTENT_TYPE,
            "authorization": f"Bearer {self.api_key}",
        }
        response = requests.post(url, json=payload, headers=headers, 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