chore: add nested llm config support to LLM reranker

This commit is contained in:
kartik-mem0
2026-03-18 22:12:04 +05:30
parent 577a5a2feb
commit 5154174342
2 changed files with 32 additions and 15 deletions
+7 -1
View File
@@ -1,4 +1,5 @@
from typing import Optional from typing import Any, Dict, Optional
from pydantic import Field from pydantic import Field
from mem0.configs.rerankers.base import BaseRerankerConfig from mem0.configs.rerankers.base import BaseRerankerConfig
@@ -46,3 +47,8 @@ class LLMRerankerConfig(BaseRerankerConfig):
default=None, default=None,
description="Custom prompt template for scoring documents" description="Custom prompt template for scoring documents"
) )
llm: Optional[Dict[str, Any]] = Field(
default=None,
description="Nested LLM configuration with 'provider' and 'config' keys. "
"Overrides top-level provider/model/api_key when provided.",
)
+25 -14
View File
@@ -1,10 +1,10 @@
import re import re
from typing import List, Dict, Any, Union from typing import Any, Dict, List, Union
from mem0.reranker.base import BaseReranker
from mem0.utils.factory import LlmFactory
from mem0.configs.rerankers.base import BaseRerankerConfig from mem0.configs.rerankers.base import BaseRerankerConfig
from mem0.configs.rerankers.llm import LLMRerankerConfig from mem0.configs.rerankers.llm import LLMRerankerConfig
from mem0.reranker.base import BaseReranker
from mem0.utils.factory import LlmFactory
class LLMReranker(BaseReranker): class LLMReranker(BaseReranker):
@@ -33,19 +33,30 @@ class LLMReranker(BaseReranker):
self.config = config self.config = config
# Create LLM configuration for the factory # If a nested ``llm`` dict is provided (e.g. for non-OpenAI providers
llm_config = { # like Ollama that need provider-specific fields such as
"model": self.config.model, # ``ollama_base_url``), use it to configure the LLM factory.
"temperature": self.config.temperature, if self.config.llm:
"max_tokens": self.config.max_tokens, nested = self.config.llm
} llm_provider = nested.get("provider", self.config.provider)
llm_config: dict = dict(nested.get("config") or {})
# Add API key if provided llm_config.setdefault("model", self.config.model)
if self.config.api_key: llm_config.setdefault("temperature", self.config.temperature)
llm_config["api_key"] = self.config.api_key llm_config.setdefault("max_tokens", self.config.max_tokens)
if self.config.api_key:
llm_config.setdefault("api_key", self.config.api_key)
else:
llm_provider = self.config.provider
llm_config = {
"model": self.config.model,
"temperature": self.config.temperature,
"max_tokens": self.config.max_tokens,
}
if self.config.api_key:
llm_config["api_key"] = self.config.api_key
# Initialize LLM using the factory # Initialize LLM using the factory
self.llm = LlmFactory.create(self.config.provider, llm_config) self.llm = LlmFactory.create(llm_provider, llm_config)
# Default scoring prompt # Default scoring prompt
self.scoring_prompt = getattr(self.config, 'scoring_prompt', None) or self._get_default_prompt() self.scoring_prompt = getattr(self.config, 'scoring_prompt', None) or self._get_default_prompt()