jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

fix: correct argument for cohere_api_key

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-06-22 21:13:44 -03:00
commit 659793fcee
1 changed files with 3 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

6
src/backend/base/langflow/components/retrievers/CohereRerank.py
View file

@ -11,7 +11,7 @@ from langflow.schema import Data
class CohereRerankComponent(LCVectorStoreComponent): class CohereRerankComponent(LCVectorStoreComponent):
display_name = "Cohere Rerank" display_name = "Cohere Rerank"
description = "Rerank documents using the Cohere API." description = "Rerank documents using the Cohere API and a retriever."
icon = "Cohere" icon = "Cohere"
inputs = [ inputs = [
@ -38,13 +38,13 @@ class CohereRerankComponent(LCVectorStoreComponent):
def build_base_retriever(self) -> Retriever: def build_base_retriever(self) -> Retriever:
cohere_reranker = CohereRerank( cohere_reranker = CohereRerank(
api_key=self.api_key, model=self.model, top_n=self.top_n, user_agent=self.user_agent cohere_api_key=self.api_key, model=self.model, top_n=self.top_n, user_agent=self.user_agent
) )
retriever = ContextualCompressionRetriever(base_compressor=cohere_reranker, base_retriever=self.retriever) retriever = ContextualCompressionRetriever(base_compressor=cohere_reranker, base_retriever=self.retriever)
return retriever return retriever
async def search_documents(self) -> List[Data]: async def search_documents(self) -> List[Data]:
retriever = self.build_base_retriever() retriever: ContextualCompressionRetriever = self.build_base_retriever()
documents = await retriever.ainvoke(self.search_query) documents = await retriever.ainvoke(self.search_query)
data = self.to_data(documents) data = self.to_data(documents)
self.status = data self.status = data