Refactored base class config for llms (#3241)
This commit is contained in:
+34
-10
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user