my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import os
#2	from typing import List, Dict, Any
#3
#4	from mem0.reranker.base import BaseReranker
#5
#6	try:
#7	import cohere
#8	COHERE_AVAILABLE = True
#9	except ImportError:
#10	COHERE_AVAILABLE = False
#11
#12
#13	class CohereReranker(BaseReranker):
#14	"""Cohere-based reranker implementation."""
#15
#16	def __init__(self, config):
#17	"""
#18	Initialize Cohere reranker.
#19
#20	Args:
#21	config: CohereRerankerConfig object with configuration parameters
#22	"""
#23	if not COHERE_AVAILABLE:
#24	raise ImportError("cohere package is required for CohereReranker. Install with: pip install cohere")
#25
#26	self.config = config
#27	self.api_key = config.api_key or os.getenv("COHERE_API_KEY")
#28	if not self.api_key:
#29	raise ValueError("Cohere API key is required. Set COHERE_API_KEY environment variable or pass api_key in config.")
#30
#31	self.model = config.model
#32	self.client = cohere.Client(self.api_key)
#33
#34	def rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
#35	"""
#36	Rerank documents using Cohere's rerank API.
#37
#38	Args:
#39	query: The search query
#40	documents: List of documents to rerank
#41	top_k: Number of top documents to return
#42
#43	Returns:
#44	List of reranked documents with rerank_score
#45	"""
#46	if not documents:
#47	return documents
#48
#49	# Extract text content for reranking
#50	doc_texts = []
#51	for doc in documents:
#52	if 'memory' in doc:
#53	doc_texts.append(doc['memory'])
#54	elif 'text' in doc:
#55	doc_texts.append(doc['text'])
#56	elif 'content' in doc:
#57	doc_texts.append(doc['content'])
#58	else:
#59	doc_texts.append(str(doc))
#60
#61	try:
#62	# Call Cohere rerank API
#63	response = self.client.rerank(
#64	model=self.model,
#65	query=query,
#66	documents=doc_texts,
#67	top_n=top_k or self.config.top_k or len(documents),
#68	return_documents=self.config.return_documents,
#69	max_chunks_per_doc=self.config.max_chunks_per_doc,
#70	)
#71
#72	# Create reranked results
#73	reranked_docs = []
#74	for result in response.results:
#75	original_doc = documents[result.index].copy()
#76	original_doc['rerank_score'] = result.relevance_score
#77	reranked_docs.append(original_doc)
#78
#79	return reranked_docs
#80
#81	except Exception:
#82	# Fallback to original order if reranking fails
#83	for doc in documents:
#84	doc['rerank_score'] = 0.0
#85	return documents[:top_k] if top_k else documents

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public