[docs]classDashScopeRerank(BaseDocumentCompressor):"""Document compressor that uses `DashScope Rerank API`."""client:Any=None"""DashScope client to use for compressing documents."""model:Optional[str]=None"""Model to use for reranking."""top_n:Optional[int]=3"""Number of documents to return."""dashscope_api_key:Optional[str]=Field(None,alias="api_key")"""DashScope API key. Must be specified directly or via environment variable DASHSCOPE_API_KEY."""model_config=ConfigDict(populate_by_name=True,arbitrary_types_allowed=True,extra="forbid",)@model_validator(mode="before")@classmethoddefvalidate_environment(cls,values:Dict)->Any:"""Validate that api key and python package exists in environment."""ifnotvalues.get("client"):try:importdashscopeexceptImportError:raiseImportError("Could not import dashscope python package. ""Please install it with `pip install dashscope`.")values["client"]=dashscope.TextReRankvalues["dashscope_api_key"]=get_from_dict_or_env(values,"dashscope_api_key","DASHSCOPE_API_KEY")values["model"]=dashscope.TextReRank.Models.gte_rerankreturnvalues
[docs]defrerank(self,documents:Sequence[Union[str,Document,dict]],query:str,*,top_n:Optional[int]=-1,)->List[Dict[str,Any]]:"""Returns an ordered list of documents ordered by their relevance to the provided query. Args: query: The query to use for reranking. documents: A sequence of documents to rerank. top_n : The number of results to return. If None returns all results. Defaults to self.top_n. """# noqa: E501iflen(documents)==0:# to avoid empty api callreturn[]docs=[doc.page_contentifisinstance(doc,Document)elsedocfordocindocuments]top_n=top_nif(top_nisNoneortop_n>0)elseself.top_nresults=self.client.call(model=self.model,query=query,documents=docs,top_n=top_n,return_documents=False,api_key=self.dashscope_api_key,)result_dicts=[]forresinresults.output.results:result_dicts.append({"index":res.index,"relevance_score":res.relevance_score,})returnresult_dicts
[docs]defcompress_documents(self,documents:Sequence[Document],query:str,callbacks:Optional[Callbacks]=None,)->Sequence[Document]:""" Compress documents using DashScope's rerank API. Args: documents: A sequence of documents to compress. query: The query to use for compressing the documents. callbacks: Callbacks to run during the compression process. Returns: A sequence of compressed documents. """compressed=[]forresinself.rerank(documents,query):doc=documents[res["index"]]doc_copy=Document(doc.page_content,metadata=deepcopy(doc.metadata))doc_copy.metadata["relevance_score"]=res["relevance_score"]compressed.append(doc_copy)returncompressed