Refactored base class config for llms (#3241)

This commit is contained in:
Parshva Daftari
2025-08-05 15:42:06 -07:00
committed by GitHub
parent f59ef3f2e2
commit 6e8c6c1cb7
25 changed files with 942 additions and 329 deletions
+34 -10
View File
@@ -1,17 +1,37 @@
import os
from typing import Dict, List, Optional
from typing import Dict, List, Optional, Union
try:
import anthropic
except ImportError:
raise ImportError("The 'anthropic' library is required. Please install it using 'pip install anthropic'.")
from mem0.configs.llms.anthropic import AnthropicConfig
from mem0.configs.llms.base import BaseLlmConfig
from mem0.llms.base import LLMBase
class AnthropicLLM(LLMBase):
def __init__(self, config: Optional[BaseLlmConfig] = None):
def __init__(self, config: Optional[Union[BaseLlmConfig, AnthropicConfig, Dict]] = None):
# Convert to AnthropicConfig if needed
if config is None:
config = AnthropicConfig()
elif isinstance(config, dict):
config = AnthropicConfig(**config)
elif isinstance(config, BaseLlmConfig) and not isinstance(config, AnthropicConfig):
# Convert BaseLlmConfig to AnthropicConfig
config = AnthropicConfig(
model=config.model,
temperature=config.temperature,
api_key=config.api_key,
max_tokens=config.max_tokens,
top_p=config.top_p,
top_k=config.top_k,
enable_vision=config.enable_vision,
vision_details=config.vision_details,
http_client_proxies=config.http_client,
)
super().__init__(config)
if not self.config.model:
@@ -26,6 +46,7 @@ class AnthropicLLM(LLMBase):
response_format=None,
tools: Optional[List[Dict]] = None,
tool_choice: str = "auto",
**kwargs,
):
"""
Generate a response based on the given messages using Anthropic.
@@ -35,6 +56,7 @@ class AnthropicLLM(LLMBase):
response_format (str or object, optional): Format of the response. Defaults to "text".
tools (list, optional): List of tools that the model can call. Defaults to None.
tool_choice (str, optional): Tool choice method. Defaults to "auto".
**kwargs: Additional Anthropic-specific parameters.
Returns:
str: The generated response.
@@ -48,14 +70,16 @@ class AnthropicLLM(LLMBase):
else:
filtered_messages.append(message)
params = {
"model": self.config.model,
"messages": filtered_messages,
"system": system_message,
"temperature": self.config.temperature,
"max_tokens": self.config.max_tokens,
"top_p": self.config.top_p,
}
# Get common parameters
params = self._get_common_params(**kwargs)
params.update(
{
"model": self.config.model,
"messages": filtered_messages,
"system": system_message,
}
)
if tools: # TODO: Remove tools if no issues found with new memory addition logic
params["tools"] = tools
params["tool_choice"] = tool_choice
+44 -16
View File
@@ -1,16 +1,36 @@
import json
import os
from typing import Dict, List, Optional
from typing import Dict, List, Optional, Union
from openai import AzureOpenAI
from mem0.configs.llms.azure import AzureOpenAIConfig
from mem0.configs.llms.base import BaseLlmConfig
from mem0.llms.base import LLMBase
from mem0.memory.utils import extract_json
class AzureOpenAILLM(LLMBase):
def __init__(self, config: Optional[BaseLlmConfig] = None):
def __init__(self, config: Optional[Union[BaseLlmConfig, AzureOpenAIConfig, Dict]] = None):
# Convert to AzureOpenAIConfig if needed
if config is None:
config = AzureOpenAIConfig()
elif isinstance(config, dict):
config = AzureOpenAIConfig(**config)
elif isinstance(config, BaseLlmConfig) and not isinstance(config, AzureOpenAIConfig):
# Convert BaseLlmConfig to AzureOpenAIConfig
config = AzureOpenAIConfig(
model=config.model,
temperature=config.temperature,
api_key=config.api_key,
max_tokens=config.max_tokens,
top_p=config.top_p,
top_k=config.top_k,
enable_vision=config.enable_vision,
vision_details=config.vision_details,
http_client_proxies=config.http_client,
)
super().__init__(config)
# Model name should match the custom deployment name chosen for it.
@@ -68,6 +88,7 @@ class AzureOpenAILLM(LLMBase):
response_format=None,
tools: Optional[List[Dict]] = None,
tool_choice: str = "auto",
**kwargs,
):
"""
Generate a response based on the given messages using Azure OpenAI.
@@ -77,6 +98,7 @@ class AzureOpenAILLM(LLMBase):
response_format (str or object, optional): Format of the response. Defaults to "text".
tools (list, optional): List of tools that the model can call. Defaults to None.
tool_choice (str, optional): Tool choice method. Defaults to "auto".
**kwargs: Additional Azure OpenAI-specific parameters.
Returns:
str: The generated response.
@@ -88,23 +110,29 @@ class AzureOpenAILLM(LLMBase):
messages[-1]["content"] = user_prompt
common_params = {
"model": self.config.model,
"messages": messages,
}
# Get common parameters
params = self._get_common_params(**kwargs)
params.update(
{
"model": self.config.model,
"messages": messages,
}
)
if self.config.model in {"o3-mini", "o1-preview", "o1"}:
params = common_params
# Use common params for these models
pass
else:
params = {
**common_params,
"temperature": self.config.temperature,
"max_tokens": self.config.max_tokens,
"top_p": self.config.top_p,
}
if response_format:
params["response_format"] = response_format
if tools: # TODO: Remove tools if no issues found with new memory addition logic
# Add additional parameters for other models
params.update(
{
"temperature": self.config.temperature,
"max_tokens": self.config.max_tokens,
"top_p": self.config.top_p,
}
)
if tools:
params["tools"] = tools
params["tool_choice"] = tool_choice
+51 -6
View File
@@ -1,23 +1,49 @@
from abc import ABC, abstractmethod
from typing import Dict, List, Optional
from typing import Dict, List, Optional, Union
from mem0.configs.llms.base import BaseLlmConfig
class LLMBase(ABC):
def __init__(self, config: Optional[BaseLlmConfig] = None):
"""
Base class for all LLM providers.
Handles common functionality and delegates provider-specific logic to subclasses.
"""
def __init__(self, config: Optional[Union[BaseLlmConfig, Dict]] = None):
"""Initialize a base LLM class
:param config: LLM configuration option class, defaults to None
:type config: Optional[BaseLlmConfig], optional
:param config: LLM configuration option class or dict, defaults to None
:type config: Optional[Union[BaseLlmConfig, Dict]], optional
"""
if config is None:
self.config = BaseLlmConfig()
elif isinstance(config, dict):
# Handle dict-based configuration (backward compatibility)
self.config = BaseLlmConfig(**config)
else:
self.config = config
# Validate configuration
self._validate_config()
def _validate_config(self):
"""
Validate the configuration.
Override in subclasses to add provider-specific validation.
"""
if not hasattr(self.config, "model"):
raise ValueError("Configuration must have a 'model' attribute")
if not hasattr(self.config, "api_key") and not hasattr(self.config, "api_key"):
# Check if API key is available via environment variable
# This will be handled by individual providers
pass
@abstractmethod
def generate_response(self, messages, tools: Optional[List[Dict]] = None, tool_choice: str = "auto"):
def generate_response(
self, messages: List[Dict[str, str]], tools: Optional[List[Dict]] = None, tool_choice: str = "auto", **kwargs
):
"""
Generate a response based on the given messages.
@@ -25,8 +51,27 @@ class LLMBase(ABC):
messages (list): List of message dicts containing 'role' and 'content'.
tools (list, optional): List of tools that the model can call. Defaults to None.
tool_choice (str, optional): Tool choice method. Defaults to "auto".
**kwargs: Additional provider-specific parameters.
Returns:
str: The generated response.
str or dict: The generated response.
"""
pass
def _get_common_params(self, **kwargs) -> Dict:
"""
Get common parameters that most providers use.
Returns:
Dict: Common parameters dictionary.
"""
params = {
"temperature": self.config.temperature,
"max_tokens": self.config.max_tokens,
"top_p": self.config.top_p,
}
# Add provider-specific parameters from kwargs
params.update(kwargs)
return params
+32 -9
View File
@@ -1,16 +1,36 @@
import json
import os
from typing import Dict, List, Optional
from typing import Dict, List, Optional, Union
from openai import OpenAI
from mem0.configs.llms.base import BaseLlmConfig
from mem0.configs.llms.deepseek import DeepSeekConfig
from mem0.llms.base import LLMBase
from mem0.memory.utils import extract_json
class DeepSeekLLM(LLMBase):
def __init__(self, config: Optional[BaseLlmConfig] = None):
def __init__(self, config: Optional[Union[BaseLlmConfig, DeepSeekConfig, Dict]] = None):
# Convert to DeepSeekConfig if needed
if config is None:
config = DeepSeekConfig()
elif isinstance(config, dict):
config = DeepSeekConfig(**config)
elif isinstance(config, BaseLlmConfig) and not isinstance(config, DeepSeekConfig):
# Convert BaseLlmConfig to DeepSeekConfig
config = DeepSeekConfig(
model=config.model,
temperature=config.temperature,
api_key=config.api_key,
max_tokens=config.max_tokens,
top_p=config.top_p,
top_k=config.top_k,
enable_vision=config.enable_vision,
vision_details=config.vision_details,
http_client_proxies=config.http_client,
)
super().__init__(config)
if not self.config.model:
@@ -56,6 +76,7 @@ class DeepSeekLLM(LLMBase):
response_format=None,
tools: Optional[List[Dict]] = None,
tool_choice: str = "auto",
**kwargs,
):
"""
Generate a response based on the given messages using DeepSeek.
@@ -65,17 +86,19 @@ class DeepSeekLLM(LLMBase):
response_format (str or object, optional): Format of the response. Defaults to "text".
tools (list, optional): List of tools that the model can call. Defaults to None.
tool_choice (str, optional): Tool choice method. Defaults to "auto".
**kwargs: Additional DeepSeek-specific parameters.
Returns:
str: The generated response.
"""
params = {
"model": self.config.model,
"messages": messages,
"temperature": self.config.temperature,
"max_tokens": self.config.max_tokens,
"top_p": self.config.top_p,
}
# Get common parameters
params = self._get_common_params(**kwargs)
params.update(
{
"model": self.config.model,
"messages": messages,
}
)
if tools:
params["tools"] = tools
+78 -14
View File
@@ -1,13 +1,35 @@
from typing import Dict, List, Optional
import json
from typing import Dict, List, Optional, Union
from openai import OpenAI
from mem0.configs.llms.base import BaseLlmConfig
from mem0.configs.llms.lmstudio import LMStudioConfig
from mem0.llms.base import LLMBase
from mem0.memory.utils import extract_json
class LMStudioLLM(LLMBase):
def __init__(self, config: Optional[BaseLlmConfig] = None):
def __init__(self, config: Optional[Union[BaseLlmConfig, LMStudioConfig, Dict]] = None):
# Convert to LMStudioConfig if needed
if config is None:
config = LMStudioConfig()
elif isinstance(config, dict):
config = LMStudioConfig(**config)
elif isinstance(config, BaseLlmConfig) and not isinstance(config, LMStudioConfig):
# Convert BaseLlmConfig to LMStudioConfig
config = LMStudioConfig(
model=config.model,
temperature=config.temperature,
api_key=config.api_key,
max_tokens=config.max_tokens,
top_p=config.top_p,
top_k=config.top_k,
enable_vision=config.enable_vision,
vision_details=config.vision_details,
http_client_proxies=config.http_client,
)
super().__init__(config)
self.config.model = (
@@ -18,12 +40,43 @@ class LMStudioLLM(LLMBase):
self.client = OpenAI(base_url=self.config.lmstudio_base_url, api_key=self.config.api_key)
def _parse_response(self, response, tools):
"""
Process the response based on whether tools are used or not.
Args:
response: The raw response from API.
tools: The list of tools provided in the request.
Returns:
str or dict: The processed response.
"""
if tools:
processed_response = {
"content": response.choices[0].message.content,
"tool_calls": [],
}
if response.choices[0].message.tool_calls:
for tool_call in response.choices[0].message.tool_calls:
processed_response["tool_calls"].append(
{
"name": tool_call.function.name,
"arguments": json.loads(extract_json(tool_call.function.arguments)),
}
)
return processed_response
else:
return response.choices[0].message.content
def generate_response(
self,
messages: List[Dict[str, str]],
response_format: dict = {"type": "json_object"},
response_format=None,
tools: Optional[List[Dict]] = None,
tool_choice: str = "auto",
**kwargs,
):
"""
Generate a response based on the given messages using LM Studio.
@@ -33,21 +86,32 @@ class LMStudioLLM(LLMBase):
response_format (str or object, optional): Format of the response. Defaults to "text".
tools (list, optional): List of tools that the model can call. Defaults to None.
tool_choice (str, optional): Tool choice method. Defaults to "auto".
**kwargs: Additional LM Studio-specific parameters.
Returns:
str: The generated response.
"""
params = {
"model": self.config.model,
"messages": messages,
"temperature": self.config.temperature,
"max_tokens": self.config.max_tokens,
"top_p": self.config.top_p,
}
if response_format:
params["response_format"] = response_format
if self.config.lmstudio_response_format is not None:
# Get common parameters
params = self._get_common_params(**kwargs)
params.update(
{
"model": self.config.model,
"messages": messages,
}
)
# Handle response format - LM Studio defaults to json_object
if self.config.lmstudio_response_format:
params["response_format"] = self.config.lmstudio_response_format
elif response_format:
params["response_format"] = response_format
else:
# Default to json_object for LM Studio
params["response_format"] = {"type": "json_object"}
if tools:
params["tools"] = tools
params["tool_choice"] = tool_choice
response = self.client.chat.completions.create(**params)
return response.choices[0].message.content
return self._parse_response(response, tools)
+44 -32
View File
@@ -1,4 +1,4 @@
from typing import Dict, List, Optional
from typing import Dict, List, Optional, Union
try:
from ollama import Client
@@ -6,25 +6,37 @@ except ImportError:
raise ImportError("The 'ollama' library is required. Please install it using 'pip install ollama'.")
from mem0.configs.llms.base import BaseLlmConfig
from mem0.configs.llms.ollama import OllamaConfig
from mem0.llms.base import LLMBase
class OllamaLLM(LLMBase):
def __init__(self, config: Optional[BaseLlmConfig] = None):
def __init__(self, config: Optional[Union[BaseLlmConfig, OllamaConfig, Dict]] = None):
# Convert to OllamaConfig if needed
if config is None:
config = OllamaConfig()
elif isinstance(config, dict):
config = OllamaConfig(**config)
elif isinstance(config, BaseLlmConfig) and not isinstance(config, OllamaConfig):
# Convert BaseLlmConfig to OllamaConfig
config = OllamaConfig(
model=config.model,
temperature=config.temperature,
api_key=config.api_key,
max_tokens=config.max_tokens,
top_p=config.top_p,
top_k=config.top_k,
enable_vision=config.enable_vision,
vision_details=config.vision_details,
http_client_proxies=config.http_client,
)
super().__init__(config)
if not self.config.model:
self.config.model = "llama3.1:70b"
self.client = Client(host=self.config.ollama_base_url)
self._ensure_model_exists()
def _ensure_model_exists(self):
"""
Ensure the specified model exists locally. If not, pull it from Ollama.
"""
local_models = self.client.list()["models"]
if not any(model.get("name") == self.config.model for model in local_models):
self.client.pull(self.config.model)
self.client = Client(host=self.config.ollama_base_url)
def _parse_response(self, response, tools):
"""
@@ -39,22 +51,18 @@ class OllamaLLM(LLMBase):
"""
if tools:
processed_response = {
"content": response["message"]["content"],
"content": response["message"]["content"] if isinstance(response, dict) else response.message.content,
"tool_calls": [],
}
if response["message"].get("tool_calls"):
for tool_call in response["message"]["tool_calls"]:
processed_response["tool_calls"].append(
{
"name": tool_call["function"]["name"],
"arguments": tool_call["function"]["arguments"],
}
)
# Ollama doesn't support tool calls in the same way, so we return the content
return processed_response
else:
return response["message"]["content"]
# Handle both dict and object responses
if isinstance(response, dict):
return response["message"]["content"]
else:
return response.message.content
def generate_response(
self,
@@ -62,33 +70,37 @@ class OllamaLLM(LLMBase):
response_format=None,
tools: Optional[List[Dict]] = None,
tool_choice: str = "auto",
**kwargs,
):
"""
Generate a response based on the given messages using OpenAI.
Generate a response based on the given messages using Ollama.
Args:
messages (list): List of message dicts containing 'role' and 'content'.
response_format (str or object, optional): Format of the response. Defaults to "text".
tools (list, optional): List of tools that the model can call. Defaults to None.
tool_choice (str, optional): Tool choice method. Defaults to "auto".
**kwargs: Additional Ollama-specific parameters.
Returns:
str: The generated response.
"""
# Build parameters for Ollama
params = {
"model": self.config.model,
"messages": messages,
"options": {
"temperature": self.config.temperature,
"num_predict": self.config.max_tokens,
"top_p": self.config.top_p,
},
}
if response_format:
params["format"] = "json"
if tools:
params["tools"] = tools
# Add options for Ollama (temperature, num_predict, top_p)
options = {
"temperature": self.config.temperature,
"num_predict": self.config.max_tokens,
"top_p": self.config.top_p,
}
params["options"] = options
# Remove OpenAI-specific parameters that Ollama doesn't support
params.pop("max_tokens", None) # Ollama uses different parameter names
response = self.client.chat(**params)
return self._parse_response(response, tools)
+33 -22
View File
@@ -1,17 +1,36 @@
import json
import os
import warnings
from typing import Dict, List, Optional
from typing import Dict, List, Optional, Union
from openai import OpenAI
from mem0.configs.llms.base import BaseLlmConfig
from mem0.configs.llms.openai import OpenAIConfig
from mem0.llms.base import LLMBase
from mem0.memory.utils import extract_json
class OpenAILLM(LLMBase):
def __init__(self, config: Optional[BaseLlmConfig] = None):
def __init__(self, config: Optional[Union[BaseLlmConfig, OpenAIConfig, Dict]] = None):
# Convert to OpenAIConfig if needed
if config is None:
config = OpenAIConfig()
elif isinstance(config, dict):
config = OpenAIConfig(**config)
elif isinstance(config, BaseLlmConfig) and not isinstance(config, OpenAIConfig):
# Convert BaseLlmConfig to OpenAIConfig
config = OpenAIConfig(
model=config.model,
temperature=config.temperature,
api_key=config.api_key,
max_tokens=config.max_tokens,
top_p=config.top_p,
top_k=config.top_k,
enable_vision=config.enable_vision,
vision_details=config.vision_details,
http_client_proxies=config.http_client,
)
super().__init__(config)
if not self.config.model:
@@ -26,18 +45,7 @@ class OpenAILLM(LLMBase):
)
else:
api_key = self.config.api_key or os.getenv("OPENAI_API_KEY")
base_url = (
self.config.openai_base_url
or os.getenv("OPENAI_API_BASE")
or os.getenv("OPENAI_BASE_URL")
or "https://api.openai.com/v1"
)
if os.environ.get("OPENAI_API_BASE"):
warnings.warn(
"The environment variable 'OPENAI_API_BASE' is deprecated and will be removed in the 0.1.80. "
"Please use 'OPENAI_BASE_URL' instead.",
DeprecationWarning,
)
base_url = self.config.openai_base_url or os.getenv("OPENAI_BASE_URL") or "https://api.openai.com/v1"
self.client = OpenAI(api_key=api_key, base_url=base_url)
@@ -77,6 +85,7 @@ class OpenAILLM(LLMBase):
response_format=None,
tools: Optional[List[Dict]] = None,
tool_choice: str = "auto",
**kwargs,
):
"""
Generate a JSON response based on the given messages using OpenAI.
@@ -86,17 +95,19 @@ class OpenAILLM(LLMBase):
response_format (str or object, optional): Format of the response. Defaults to "text".
tools (list, optional): List of tools that the model can call. Defaults to None.
tool_choice (str, optional): Tool choice method. Defaults to "auto".
**kwargs: Additional OpenAI-specific parameters.
Returns:
json: The generated response.
"""
params = {
"model": self.config.model,
"messages": messages,
"temperature": self.config.temperature,
"max_tokens": self.config.max_tokens,
"top_p": self.config.top_p,
}
# Get common parameters
params = self._get_common_params(**kwargs)
params.update(
{
"model": self.config.model,
"messages": messages,
}
)
if os.getenv("OPENROUTER_API_KEY"):
openrouter_params = {}
+33 -14
View File
@@ -1,16 +1,36 @@
import json
import os
from typing import Dict, List, Optional
from typing import Dict, List, Optional, Union
from openai import OpenAI
from mem0.configs.llms.base import BaseLlmConfig
from mem0.configs.llms.vllm import VllmConfig
from mem0.llms.base import LLMBase
from mem0.memory.utils import extract_json
class VllmLLM(LLMBase):
def __init__(self, config: Optional[BaseLlmConfig] = None):
def __init__(self, config: Optional[Union[BaseLlmConfig, VllmConfig, Dict]] = None):
# Convert to VllmConfig if needed
if config is None:
config = VllmConfig()
elif isinstance(config, dict):
config = VllmConfig(**config)
elif isinstance(config, BaseLlmConfig) and not isinstance(config, VllmConfig):
# Convert BaseLlmConfig to VllmConfig
config = VllmConfig(
model=config.model,
temperature=config.temperature,
api_key=config.api_key,
max_tokens=config.max_tokens,
top_p=config.top_p,
top_k=config.top_k,
enable_vision=config.enable_vision,
vision_details=config.vision_details,
http_client_proxies=config.http_client,
)
super().__init__(config)
if not self.config.model:
@@ -18,8 +38,7 @@ class VllmLLM(LLMBase):
self.config.api_key = self.config.api_key or os.getenv("VLLM_API_KEY") or "vllm-api-key"
base_url = self.config.vllm_base_url or os.getenv("VLLM_BASE_URL")
self.client = OpenAI(base_url=base_url, api_key=self.config.api_key)
self.client = OpenAI(api_key=self.config.api_key, base_url=base_url)
def _parse_response(self, response, tools):
"""
@@ -57,6 +76,7 @@ class VllmLLM(LLMBase):
response_format=None,
tools: Optional[List[Dict]] = None,
tool_choice: str = "auto",
**kwargs,
):
"""
Generate a response based on the given messages using vLLM.
@@ -66,20 +86,19 @@ class VllmLLM(LLMBase):
response_format (str or object, optional): Format of the response. Defaults to "text".
tools (list, optional): List of tools that the model can call. Defaults to None.
tool_choice (str, optional): Tool choice method. Defaults to "auto".
**kwargs: Additional vLLM-specific parameters.
Returns:
str: The generated response.
"""
params = {
"model": self.config.model,
"messages": messages,
"temperature": self.config.temperature,
"max_tokens": self.config.max_tokens,
"top_p": self.config.top_p,
}
if response_format:
params["response_format"] = response_format
# Get common parameters
params = self._get_common_params(**kwargs)
params.update(
{
"model": self.config.model,
"messages": messages,
}
)
if tools:
params["tools"] = tools