Remove tools from LLMs (#2363)

This commit is contained in:
Anusha Yella
2025-03-14 17:42:48 +05:30
committed by GitHub
parent 4be426f762
commit ee80a43810
21 changed files with 418 additions and 1071 deletions

View File

@@ -4,70 +4,50 @@ from typing import Dict, List, Optional
try:
import litellm
except ImportError:
raise ImportError("The 'litellm' library is required. Please install it using 'pip install litellm'.")
raise ImportError(
"The 'litellm' library is required. Please install it using 'pip install litellm'."
)
from mem0.configs.llms.base import BaseLlmConfig
from mem0.llms.base import LLMBase
class LiteLLM(LLMBase):
"""
A class for interacting with LiteLLM's language models using the specified configuration.
"""
def __init__(self, config: Optional[BaseLlmConfig] = None):
"""
Initializes the LiteLLM instance with the given configuration.
Args:
config (Optional[BaseLlmConfig]): Configuration settings for the language model.
"""
super().__init__(config)
if not self.config.model:
self.config.model = "gpt-4o-mini"
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(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=None,
tools: Optional[List[Dict]] = None,
tool_choice: str = "auto",
):
response_format: Optional[str] = None,
) -> str:
"""
Generate a response based on the given messages using Litellm.
Generates a response using LiteLLM based on the provided messages.
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".
messages (List[Dict[str, str]]): A list of dictionaries, each containing a 'role' and 'content' key.
response_format (Optional[str]): The desired format of the response. Defaults to None.
Returns:
str: The generated response.
str: The generated response from the model.
"""
if not litellm.supports_function_calling(self.config.model):
raise ValueError(f"Model '{self.config.model}' in litellm does not support function calling.")
raise ValueError(
f"Model '{self.config.model}' in LiteLLM does not support function calling."
)
params = {
"model": self.config.model,
@@ -78,9 +58,6 @@ class LiteLLM(LLMBase):
}
if response_format:
params["response_format"] = response_format
if tools: # TODO: Remove tools if no issues found with new memory addition logic
params["tools"] = tools
params["tool_choice"] = tool_choice
response = litellm.completion(**params)
return self._parse_response(response, tools)
return response.choices[0].message.content