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