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