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
+56
View File
@@ -0,0 +1,56 @@
from typing import Optional
from mem0.configs.llms.base import BaseLlmConfig
class AnthropicConfig(BaseLlmConfig):
"""
Configuration class for Anthropic-specific parameters.
Inherits from BaseLlmConfig and adds Anthropic-specific settings.
"""
def __init__(
self,
# Base parameters
model: Optional[str] = None,
temperature: float = 0.1,
api_key: Optional[str] = None,
max_tokens: int = 2000,
top_p: float = 0.1,
top_k: int = 1,
enable_vision: bool = False,
vision_details: Optional[str] = "auto",
http_client_proxies: Optional[dict] = None,
# Anthropic-specific parameters
anthropic_base_url: Optional[str] = None,
):
"""
Initialize Anthropic configuration.
Args:
model: Anthropic model to use, defaults to None
temperature: Controls randomness, defaults to 0.1
api_key: Anthropic API key, defaults to None
max_tokens: Maximum tokens to generate, defaults to 2000
top_p: Nucleus sampling parameter, defaults to 0.1
top_k: Top-k sampling parameter, defaults to 1
enable_vision: Enable vision capabilities, defaults to False
vision_details: Vision detail level, defaults to "auto"
http_client_proxies: HTTP client proxy settings, defaults to None
anthropic_base_url: Anthropic API base URL, defaults to None
"""
# Initialize base parameters
super().__init__(
model=model,
temperature=temperature,
api_key=api_key,
max_tokens=max_tokens,
top_p=top_p,
top_k=top_k,
enable_vision=enable_vision,
vision_details=vision_details,
http_client_proxies=http_client_proxies,
)
# Anthropic-specific parameters
self.anthropic_base_url = anthropic_base_url
+57
View File
@@ -0,0 +1,57 @@
from typing import Any, Dict, Optional
from mem0.configs.base import AzureConfig
from mem0.configs.llms.base import BaseLlmConfig
class AzureOpenAIConfig(BaseLlmConfig):
"""
Configuration class for Azure OpenAI-specific parameters.
Inherits from BaseLlmConfig and adds Azure OpenAI-specific settings.
"""
def __init__(
self,
# Base parameters
model: Optional[str] = None,
temperature: float = 0.1,
api_key: Optional[str] = None,
max_tokens: int = 2000,
top_p: float = 0.1,
top_k: int = 1,
enable_vision: bool = False,
vision_details: Optional[str] = "auto",
http_client_proxies: Optional[dict] = None,
# Azure OpenAI-specific parameters
azure_kwargs: Optional[Dict[str, Any]] = None,
):
"""
Initialize Azure OpenAI configuration.
Args:
model: Azure OpenAI model to use, defaults to None
temperature: Controls randomness, defaults to 0.1
api_key: Azure OpenAI API key, defaults to None
max_tokens: Maximum tokens to generate, defaults to 2000
top_p: Nucleus sampling parameter, defaults to 0.1
top_k: Top-k sampling parameter, defaults to 1
enable_vision: Enable vision capabilities, defaults to False
vision_details: Vision detail level, defaults to "auto"
http_client_proxies: HTTP client proxy settings, defaults to None
azure_kwargs: Azure-specific configuration, defaults to None
"""
# Initialize base parameters
super().__init__(
model=model,
temperature=temperature,
api_key=api_key,
max_tokens=max_tokens,
top_p=top_p,
top_k=top_k,
enable_vision=enable_vision,
vision_details=vision_details,
http_client_proxies=http_client_proxies,
)
# Azure OpenAI-specific parameters
self.azure_kwargs = AzureConfig(**(azure_kwargs or {}))
+28 -118
View File
@@ -3,12 +3,14 @@ from typing import Dict, Optional, Union
import httpx
from mem0.configs.base import AzureConfig
class BaseLlmConfig(ABC):
"""
Config for LLMs.
Base configuration for LLMs with only common parameters.
Provider-specific configurations should be handled by separate config classes.
This class contains only the parameters that are common across all LLM providers.
For provider-specific parameters, use the appropriate provider config class.
"""
def __init__(
@@ -21,89 +23,34 @@ class BaseLlmConfig(ABC):
top_k: int = 1,
enable_vision: bool = False,
vision_details: Optional[str] = "auto",
# Openrouter specific
models: Optional[list[str]] = None,
route: Optional[str] = "fallback",
openrouter_base_url: Optional[str] = None,
# Openai specific
openai_base_url: Optional[str] = None,
site_url: Optional[str] = None,
app_name: Optional[str] = None,
# Ollama specific
ollama_base_url: Optional[str] = None,
# AzureOpenAI specific
azure_kwargs: Optional[AzureConfig] = {},
# AzureOpenAI specific
http_client_proxies: Optional[Union[Dict, str]] = None,
# DeepSeek specific
deepseek_base_url: Optional[str] = None,
# XAI specific
xai_base_url: Optional[str] = None,
# Sarvam specific
sarvam_base_url: Optional[str] = "https://api.sarvam.ai/v1",
# LM Studio specific
lmstudio_base_url: Optional[str] = "http://localhost:1234/v1",
lmstudio_response_format: dict = None,
# vLLM specific
vllm_base_url: Optional[str] = "http://localhost:8000/v1",
# AWS Bedrock specific
aws_access_key_id: Optional[str] = None,
aws_secret_access_key: Optional[str] = None,
aws_region: Optional[str] = "us-west-2",
):
"""
Initializes a configuration class instance for the LLM.
Initialize a base configuration class instance for the LLM.
:param model: Controls the OpenAI model used, defaults to None
:type model: Optional[str], optional
:param temperature: Controls the randomness of the model's output.
Higher values (closer to 1) make output more random, lower values make it more deterministic, defaults to 0
:type temperature: float, optional
:param api_key: OpenAI API key to be use, defaults to None
:type api_key: Optional[str], optional
:param max_tokens: Controls how many tokens are generated, defaults to 2000
:type max_tokens: int, optional
:param top_p: Controls the diversity of words. Higher values (closer to 1) make word selection more diverse,
defaults to 1
:type top_p: float, optional
:param top_k: Controls the diversity of words. Higher values make word selection more diverse, defaults to 0
:type top_k: int, optional
:param enable_vision: Enable vision for the LLM, defaults to False
:type enable_vision: bool, optional
:param vision_details: Details of the vision to be used [low, high, auto], defaults to "auto"
:type vision_details: Optional[str], optional
:param models: Openrouter models to use, defaults to None
:type models: Optional[list[str]], optional
:param route: Openrouter route to be used, defaults to "fallback"
:type route: Optional[str], optional
:param openrouter_base_url: Openrouter base URL to be use, defaults to "https://openrouter.ai/api/v1"
:type openrouter_base_url: Optional[str], optional
:param site_url: Openrouter site URL to use, defaults to None
:type site_url: Optional[str], optional
:param app_name: Openrouter app name to use, defaults to None
:type app_name: Optional[str], optional
:param ollama_base_url: The base URL of the LLM, defaults to None
:type ollama_base_url: Optional[str], optional
:param openai_base_url: Openai base URL to be use, defaults to "https://api.openai.com/v1"
:type openai_base_url: Optional[str], optional
:param azure_kwargs: key-value arguments for the AzureOpenAI LLM model, defaults a dict inside init
:type azure_kwargs: Optional[Dict[str, Any]], defaults a dict inside init
:param http_client_proxies: The proxy server(s) settings used to create self.http_client, defaults to None
:type http_client_proxies: Optional[Dict | str], optional
:param deepseek_base_url: DeepSeek base URL to be use, defaults to None
:type deepseek_base_url: Optional[str], optional
:param xai_base_url: XAI base URL to be use, defaults to None
:type xai_base_url: Optional[str], optional
:param sarvam_base_url: Sarvam base URL to be use, defaults to "https://api.sarvam.ai/v1"
:type sarvam_base_url: Optional[str], optional
:param lmstudio_base_url: LM Studio base URL to be use, defaults to "http://localhost:1234/v1"
:type lmstudio_base_url: Optional[str], optional
:param lmstudio_response_format: LM Studio response format to be use, defaults to None
:type lmstudio_response_format: Optional[Dict], optional
:param vllm_base_url: vLLM base URL to be use, defaults to "http://localhost:8000/v1"
:type vllm_base_url: Optional[str], optional
Args:
model: The model identifier to use (e.g., "gpt-4o-mini", "claude-3-5-sonnet-20240620")
Defaults to None (will be set by provider-specific configs)
temperature: Controls the randomness of the model's output.
Higher values (closer to 1) make output more random, lower values make it more deterministic.
Range: 0.0 to 2.0. Defaults to 0.1
api_key: API key for the LLM provider. If None, will try to get from environment variables.
Defaults to None
max_tokens: Maximum number of tokens to generate in the response.
Range: 1 to 4096 (varies by model). Defaults to 2000
top_p: Nucleus sampling parameter. Controls diversity via nucleus sampling.
Higher values (closer to 1) make word selection more diverse.
Range: 0.0 to 1.0. Defaults to 0.1
top_k: Top-k sampling parameter. Limits the number of tokens considered for each step.
Higher values make word selection more diverse.
Range: 1 to 40. Defaults to 1
enable_vision: Whether to enable vision capabilities for the model.
Only applicable to vision-enabled models. Defaults to False
vision_details: Level of detail for vision processing.
Options: "low", "high", "auto". Defaults to "auto"
http_client_proxies: Proxy settings for HTTP client.
Can be a dict or string. Defaults to None
"""
self.model = model
self.temperature = temperature
self.api_key = api_key
@@ -112,41 +59,4 @@ class BaseLlmConfig(ABC):
self.top_k = top_k
self.enable_vision = enable_vision
self.vision_details = vision_details
# AzureOpenAI specific
self.http_client = httpx.Client(proxies=http_client_proxies) if http_client_proxies else None
# Openrouter specific
self.models = models
self.route = route
self.openrouter_base_url = openrouter_base_url
self.openai_base_url = openai_base_url
self.site_url = site_url
self.app_name = app_name
# Ollama specific
self.ollama_base_url = ollama_base_url
# DeepSeek specific
self.deepseek_base_url = deepseek_base_url
# AzureOpenAI specific
self.azure_kwargs = AzureConfig(**azure_kwargs) or {}
# XAI specific
self.xai_base_url = xai_base_url
# Sarvam specific
self.sarvam_base_url = sarvam_base_url
# LM Studio specific
self.lmstudio_base_url = lmstudio_base_url
self.lmstudio_response_format = lmstudio_response_format
# vLLM specific
self.vllm_base_url = vllm_base_url
# AWS Bedrock specific
self.aws_access_key_id = aws_access_key_id
self.aws_secret_access_key = aws_secret_access_key
self.aws_region = aws_region
+56
View File
@@ -0,0 +1,56 @@
from typing import Optional
from mem0.configs.llms.base import BaseLlmConfig
class DeepSeekConfig(BaseLlmConfig):
"""
Configuration class for DeepSeek-specific parameters.
Inherits from BaseLlmConfig and adds DeepSeek-specific settings.
"""
def __init__(
self,
# Base parameters
model: Optional[str] = None,
temperature: float = 0.1,
api_key: Optional[str] = None,
max_tokens: int = 2000,
top_p: float = 0.1,
top_k: int = 1,
enable_vision: bool = False,
vision_details: Optional[str] = "auto",
http_client_proxies: Optional[dict] = None,
# DeepSeek-specific parameters
deepseek_base_url: Optional[str] = None,
):
"""
Initialize DeepSeek configuration.
Args:
model: DeepSeek model to use, defaults to None
temperature: Controls randomness, defaults to 0.1
api_key: DeepSeek API key, defaults to None
max_tokens: Maximum tokens to generate, defaults to 2000
top_p: Nucleus sampling parameter, defaults to 0.1
top_k: Top-k sampling parameter, defaults to 1
enable_vision: Enable vision capabilities, defaults to False
vision_details: Vision detail level, defaults to "auto"
http_client_proxies: HTTP client proxy settings, defaults to None
deepseek_base_url: DeepSeek API base URL, defaults to None
"""
# Initialize base parameters
super().__init__(
model=model,
temperature=temperature,
api_key=api_key,
max_tokens=max_tokens,
top_p=top_p,
top_k=top_k,
enable_vision=enable_vision,
vision_details=vision_details,
http_client_proxies=http_client_proxies,
)
# DeepSeek-specific parameters
self.deepseek_base_url = deepseek_base_url
+59
View File
@@ -0,0 +1,59 @@
from typing import Any, Dict, Optional
from mem0.configs.llms.base import BaseLlmConfig
class LMStudioConfig(BaseLlmConfig):
"""
Configuration class for LM Studio-specific parameters.
Inherits from BaseLlmConfig and adds LM Studio-specific settings.
"""
def __init__(
self,
# Base parameters
model: Optional[str] = None,
temperature: float = 0.1,
api_key: Optional[str] = None,
max_tokens: int = 2000,
top_p: float = 0.1,
top_k: int = 1,
enable_vision: bool = False,
vision_details: Optional[str] = "auto",
http_client_proxies: Optional[dict] = None,
# LM Studio-specific parameters
lmstudio_base_url: Optional[str] = None,
lmstudio_response_format: Optional[Dict[str, Any]] = None,
):
"""
Initialize LM Studio configuration.
Args:
model: LM Studio model to use, defaults to None
temperature: Controls randomness, defaults to 0.1
api_key: LM Studio API key, defaults to None
max_tokens: Maximum tokens to generate, defaults to 2000
top_p: Nucleus sampling parameter, defaults to 0.1
top_k: Top-k sampling parameter, defaults to 1
enable_vision: Enable vision capabilities, defaults to False
vision_details: Vision detail level, defaults to "auto"
http_client_proxies: HTTP client proxy settings, defaults to None
lmstudio_base_url: LM Studio base URL, defaults to None
lmstudio_response_format: LM Studio response format, defaults to None
"""
# Initialize base parameters
super().__init__(
model=model,
temperature=temperature,
api_key=api_key,
max_tokens=max_tokens,
top_p=top_p,
top_k=top_k,
enable_vision=enable_vision,
vision_details=vision_details,
http_client_proxies=http_client_proxies,
)
# LM Studio-specific parameters
self.lmstudio_base_url = lmstudio_base_url or "http://localhost:1234/v1"
self.lmstudio_response_format = lmstudio_response_format
+56
View File
@@ -0,0 +1,56 @@
from typing import Optional
from mem0.configs.llms.base import BaseLlmConfig
class OllamaConfig(BaseLlmConfig):
"""
Configuration class for Ollama-specific parameters.
Inherits from BaseLlmConfig and adds Ollama-specific settings.
"""
def __init__(
self,
# Base parameters
model: Optional[str] = None,
temperature: float = 0.1,
api_key: Optional[str] = None,
max_tokens: int = 2000,
top_p: float = 0.1,
top_k: int = 1,
enable_vision: bool = False,
vision_details: Optional[str] = "auto",
http_client_proxies: Optional[dict] = None,
# Ollama-specific parameters
ollama_base_url: Optional[str] = None,
):
"""
Initialize Ollama configuration.
Args:
model: Ollama model to use, defaults to None
temperature: Controls randomness, defaults to 0.1
api_key: Ollama API key, defaults to None
max_tokens: Maximum tokens to generate, defaults to 2000
top_p: Nucleus sampling parameter, defaults to 0.1
top_k: Top-k sampling parameter, defaults to 1
enable_vision: Enable vision capabilities, defaults to False
vision_details: Vision detail level, defaults to "auto"
http_client_proxies: HTTP client proxy settings, defaults to None
ollama_base_url: Ollama base URL, defaults to None
"""
# Initialize base parameters
super().__init__(
model=model,
temperature=temperature,
api_key=api_key,
max_tokens=max_tokens,
top_p=top_p,
top_k=top_k,
enable_vision=enable_vision,
vision_details=vision_details,
http_client_proxies=http_client_proxies,
)
# Ollama-specific parameters
self.ollama_base_url = ollama_base_url
+71
View File
@@ -0,0 +1,71 @@
from typing import List, Optional
from mem0.configs.llms.base import BaseLlmConfig
class OpenAIConfig(BaseLlmConfig):
"""
Configuration class for OpenAI and OpenRouter-specific parameters.
Inherits from BaseLlmConfig and adds OpenAI-specific settings.
"""
def __init__(
self,
# Base parameters
model: Optional[str] = None,
temperature: float = 0.1,
api_key: Optional[str] = None,
max_tokens: int = 2000,
top_p: float = 0.1,
top_k: int = 1,
enable_vision: bool = False,
vision_details: Optional[str] = "auto",
http_client_proxies: Optional[dict] = None,
# OpenAI-specific parameters
openai_base_url: Optional[str] = None,
models: Optional[List[str]] = None,
route: Optional[str] = "fallback",
openrouter_base_url: Optional[str] = None,
site_url: Optional[str] = None,
app_name: Optional[str] = None,
):
"""
Initialize OpenAI configuration.
Args:
model: OpenAI model to use, defaults to None
temperature: Controls randomness, defaults to 0.1
api_key: OpenAI API key, defaults to None
max_tokens: Maximum tokens to generate, defaults to 2000
top_p: Nucleus sampling parameter, defaults to 0.1
top_k: Top-k sampling parameter, defaults to 1
enable_vision: Enable vision capabilities, defaults to False
vision_details: Vision detail level, defaults to "auto"
http_client_proxies: HTTP client proxy settings, defaults to None
openai_base_url: OpenAI API base URL, defaults to None
models: List of models for OpenRouter, defaults to None
route: OpenRouter route strategy, defaults to "fallback"
openrouter_base_url: OpenRouter base URL, defaults to None
site_url: Site URL for OpenRouter, defaults to None
app_name: Application name for OpenRouter, defaults to None
"""
# Initialize base parameters
super().__init__(
model=model,
temperature=temperature,
api_key=api_key,
max_tokens=max_tokens,
top_p=top_p,
top_k=top_k,
enable_vision=enable_vision,
vision_details=vision_details,
http_client_proxies=http_client_proxies,
)
# OpenAI-specific parameters
self.openai_base_url = openai_base_url
self.models = models
self.route = route
self.openrouter_base_url = openrouter_base_url
self.site_url = site_url
self.app_name = app_name
+56
View File
@@ -0,0 +1,56 @@
from typing import Optional
from mem0.configs.llms.base import BaseLlmConfig
class VllmConfig(BaseLlmConfig):
"""
Configuration class for vLLM-specific parameters.
Inherits from BaseLlmConfig and adds vLLM-specific settings.
"""
def __init__(
self,
# Base parameters
model: Optional[str] = None,
temperature: float = 0.1,
api_key: Optional[str] = None,
max_tokens: int = 2000,
top_p: float = 0.1,
top_k: int = 1,
enable_vision: bool = False,
vision_details: Optional[str] = "auto",
http_client_proxies: Optional[dict] = None,
# vLLM-specific parameters
vllm_base_url: Optional[str] = None,
):
"""
Initialize vLLM configuration.
Args:
model: vLLM model to use, defaults to None
temperature: Controls randomness, defaults to 0.1
api_key: vLLM API key, defaults to None
max_tokens: Maximum tokens to generate, defaults to 2000
top_p: Nucleus sampling parameter, defaults to 0.1
top_k: Top-k sampling parameter, defaults to 1
enable_vision: Enable vision capabilities, defaults to False
vision_details: Vision detail level, defaults to "auto"
http_client_proxies: HTTP client proxy settings, defaults to None
vllm_base_url: vLLM base URL, defaults to None
"""
# Initialize base parameters
super().__init__(
model=model,
temperature=temperature,
api_key=api_key,
max_tokens=max_tokens,
top_p=top_p,
top_k=top_k,
enable_vision=enable_vision,
vision_details=vision_details,
http_client_proxies=http_client_proxies,
)
# vLLM-specific parameters
self.vllm_base_url = vllm_base_url or "http://localhost:8000/v1"