my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import re
#2	from typing import List, Dict, Any, Union
#3
#4	from mem0.reranker.base import BaseReranker
#5	from mem0.utils.factory import LlmFactory
#6	from mem0.configs.rerankers.base import BaseRerankerConfig
#7	from mem0.configs.rerankers.llm import LLMRerankerConfig
#8
#9
#10	class LLMReranker(BaseReranker):
#11	"""LLM-based reranker implementation."""
#12
#13	def __init__(self, config: Union[BaseRerankerConfig, LLMRerankerConfig, Dict]):
#14	"""
#15	Initialize LLM reranker.
#16
#17	Args:
#18	config: Configuration object with reranker parameters
#19	"""
#20	# Convert to LLMRerankerConfig if needed
#21	if isinstance(config, dict):
#22	config = LLMRerankerConfig(**config)
#23	elif isinstance(config, BaseRerankerConfig) and not isinstance(config, LLMRerankerConfig):
#24	# Convert BaseRerankerConfig to LLMRerankerConfig with defaults
#25	config = LLMRerankerConfig(
#26	provider=getattr(config, 'provider', 'openai'),
#27	model=getattr(config, 'model', 'gpt-4o-mini'),
#28	api_key=getattr(config, 'api_key', None),
#29	top_k=getattr(config, 'top_k', None),
#30	temperature=0.0, # Default for reranking
#31	max_tokens=100, # Default for reranking
#32	)
#33
#34	self.config = config
#35
#36	# Create LLM configuration for the factory
#37	llm_config = {
#38	"model": self.config.model,
#39	"temperature": self.config.temperature,
#40	"max_tokens": self.config.max_tokens,
#41	}
#42
#43	# Add API key if provided
#44	if self.config.api_key:
#45	llm_config["api_key"] = self.config.api_key
#46
#47	# Initialize LLM using the factory
#48	self.llm = LlmFactory.create(self.config.provider, llm_config)
#49
#50	# Default scoring prompt
#51	self.scoring_prompt = getattr(self.config, 'scoring_prompt', None) or self._get_default_prompt()
#52
#53	def _get_default_prompt(self) -> str:
#54	"""Get the default scoring prompt template."""
#55	return """You are a relevance scoring assistant. Given a query and a document, you need to score how relevant the document is to the query.
#56
#57	Score the relevance on a scale from 0.0 to 1.0, where:
#58	- 1.0 = Perfectly relevant and directly answers the query
#59	- 0.8-0.9 = Highly relevant with good information
#60	- 0.6-0.7 = Moderately relevant with some useful information
#61	- 0.4-0.5 = Slightly relevant with limited useful information
#62	- 0.0-0.3 = Not relevant or no useful information
#63
#64	Query: "{query}"
#65	Document: "{document}"
#66
#67	Provide only a single numerical score between 0.0 and 1.0. Do not include any explanation or additional text."""
#68
#69	def _extract_score(self, response_text: str) -> float:
#70	"""Extract numerical score from LLM response."""
#71	# Look for decimal numbers between 0.0 and 1.0
#72	pattern = r'\b([01](?:\.\d+)?)\b'
#73	matches = re.findall(pattern, response_text)
#74
#75	if matches:
#76	score = float(matches[0])
#77	return min(max(score, 0.0), 1.0) # Clamp between 0.0 and 1.0
#78
#79	# Fallback: return 0.5 if no valid score found
#80	return 0.5
#81
#82	def rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
#83	"""
#84	Rerank documents using LLM scoring.
#85
#86	Args:
#87	query: The search query
#88	documents: List of documents to rerank
#89	top_k: Number of top documents to return
#90
#91	Returns:
#92	List of reranked documents with rerank_score
#93	"""
#94	if not documents:
#95	return documents
#96
#97	scored_docs = []
#98
#99	for doc in documents:
#100	# Extract text content
#101	if 'memory' in doc:
#102	doc_text = doc['memory']
#103	elif 'text' in doc:
#104	doc_text = doc['text']
#105	elif 'content' in doc:
#106	doc_text = doc['content']
#107	else:
#108	doc_text = str(doc)
#109
#110	try:
#111	# Generate scoring prompt
#112	prompt = self.scoring_prompt.format(query=query, document=doc_text)
#113
#114	# Get LLM response
#115	response = self.llm.generate_response(
#116	messages=[{"role": "user", "content": prompt}]
#117	)
#118
#119	# Extract score from response
#120	score = self._extract_score(response)
#121
#122	# Create scored document
#123	scored_doc = doc.copy()
#124	scored_doc['rerank_score'] = score
#125	scored_docs.append(scored_doc)
#126
#127	except Exception:
#128	# Fallback: assign neutral score if scoring fails
#129	scored_doc = doc.copy()
#130	scored_doc['rerank_score'] = 0.5
#131	scored_docs.append(scored_doc)
#132
#133	# Sort by relevance score in descending order
#134	scored_docs.sort(key=lambda x: x['rerank_score'], reverse=True)
#135
#136	# Apply top_k limit
#137	if top_k:
#138	scored_docs = scored_docs[:top_k]
#139	elif self.config.top_k:
#140	scored_docs = scored_docs[:self.config.top_k]
#141
#142	return scored_docs

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public