github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-07-15 17:33:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

community[patch]: give reranker default client val (#17289)

Browse Source
This commit is contained in:
Bagatur 2024-02-12 17:21:53 -08:00 committed by GitHub
parent 841e5f514e
commit 22638e5927
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 15 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
libs/langchain/langchain/retrievers/document_compressors/cohere_rerank.py
View File

@ -14,14 +14,15 @@ from langchain.utils import get_from_dict_or_env
class CohereRerank(BaseDocumentCompressor): class CohereRerank(BaseDocumentCompressor):
"""Document compressor that uses `Cohere Rerank API`.""" """Document compressor that uses `Cohere Rerank API`."""
client: Any client: Any = None
"""Cohere client to use for compressing documents.""" """Cohere client to use for compressing documents."""
top_n: Optional[int] = 3 top_n: Optional[int] = 3
"""Number of documents to return.""" """Number of documents to return."""
model: str = "rerank-english-v2.0" model: str = "rerank-english-v2.0"
"""Model to use for reranking.""" """Model to use for reranking."""
cohere_api_key: Optional[str] = None cohere_api_key: Optional[str] = None
"""Cohere API key. Must be specified directly or via environment variable
COHERE_API_KEY."""
user_agent: str = "langchain" user_agent: str = "langchain"
"""Identifier for the application making the request.""" """Identifier for the application making the request."""
@ -34,18 +35,19 @@ class CohereRerank(BaseDocumentCompressor):
@root_validator(pre=True) @root_validator(pre=True)
def validate_environment(cls, values: Dict) -> Dict: def validate_environment(cls, values: Dict) -> Dict:
"""Validate that api key and python package exists in environment.""" """Validate that api key and python package exists in environment."""
try: if not values.get("client"):
import cohere try:
except ImportError: import cohere
raise ImportError( except ImportError:
"Could not import cohere python package. " raise ImportError(
"Please install it with `pip install cohere`." "Could not import cohere python package. "
"Please install it with `pip install cohere`."
)
cohere_api_key = get_from_dict_or_env(
values, "cohere_api_key", "COHERE_API_KEY"
) )
cohere_api_key = get_from_dict_or_env( client_name = values.get("user_agent", "langchain")
values, "cohere_api_key", "COHERE_API_KEY" values["client"] = cohere.Client(cohere_api_key, client_name=client_name)
)
client_name = values.get("user_agent", "langchain")
values["client"] = cohere.Client(cohere_api_key, client_name=client_name)
return values return values
def rerank( def rerank(