tools fix and formatting (#2441)

mem0/llms/azure_openai_structured.py

@@ -1,4 +1,3 @@
-import json
 import os
 from typing import Dict, List, Optional
 
@@ -36,6 +35,8 @@ class AzureOpenAIStructuredLLM(LLMBase):
         self,
         messages: List[Dict[str, str]],
         response_format: Optional[str] = None,
+        tools: Optional[List[Dict]] = None,
+        tool_choice: str = "auto",
     ) -> str:
         """
         Generate a response based on the given messages using Azure OpenAI.

mem0/llms/base.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import Dict, List, Optional
 
 from mem0.configs.llms.base import BaseLlmConfig
 
@@ -17,12 +17,14 @@ class LLMBase(ABC):
             self.config = config
 
     @abstractmethod
-    def generate_response(self, messages):
+    def generate_response(self, messages, tools: Optional[List[Dict]] = None, tool_choice: str = "auto"):
         """
         Generate a response based on the given messages.
 
         Args:
             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".
 
         Returns:
             str: The generated response.

mem0/llms/groq.py

@@ -84,4 +84,4 @@ class GroqLLM(LLMBase):
             params["tool_choice"] = tool_choice
 
         response = self.client.chat.completions.create(**params)
-        return self._parse_response(response, tools)
+        return self._parse_response(response, tools)

mem0/llms/lmstudio.py

@@ -10,7 +10,10 @@ class LMStudioLLM(LLMBase):
     def __init__(self, config: Optional[BaseLlmConfig] = None):
         super().__init__(config)
 
-        self.config.model = self.config.model or "lmstudio-community/Meta-Llama-3.1-70B-Instruct-GGUF/Meta-Llama-3.1-70B-Instruct-IQ2_M.gguf"
+        self.config.model = (
+            self.config.model
+            or "lmstudio-community/Meta-Llama-3.1-70B-Instruct-GGUF/Meta-Llama-3.1-70B-Instruct-IQ2_M.gguf"
+        )
         self.config.api_key = self.config.api_key or "lm-studio"
 
         self.client = OpenAI(base_url=self.config.lmstudio_base_url, api_key=self.config.api_key)
@@ -20,7 +23,7 @@ class LMStudioLLM(LLMBase):
         messages: List[Dict[str, str]],
         response_format: dict = {"type": "json_object"},
         tools: Optional[List[Dict]] = None,
-        tool_choice: str = "auto"
+        tool_choice: str = "auto",
     ):
         """
         Generate a response based on the given messages using LM Studio.
@@ -39,10 +42,10 @@ class LMStudioLLM(LLMBase):
             "messages": messages,
             "temperature": self.config.temperature,
             "max_tokens": self.config.max_tokens,
-            "top_p": self.config.top_p
+            "top_p": self.config.top_p,
         }
         if response_format:
             params["response_format"] = response_format
 
         response = self.client.chat.completions.create(**params)
-        return response.choices[0].message.content
+        return response.choices[0].message.content

mem0/llms/openai.py

@@ -1,3 +1,4 @@
+import json
 import os
 import warnings
 from typing import Dict, List, Optional
@@ -34,7 +35,7 @@ class OpenAILLM(LLMBase):
                 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
+                    DeprecationWarning,
                 )
 
             self.client = OpenAI(api_key=api_key, base_url=base_url)

mem0/llms/openai_structured.py

@@ -1,4 +1,3 @@
-import json
 import os
 from typing import Dict, List, Optional
 
@@ -23,6 +22,8 @@ class OpenAIStructuredLLM(LLMBase):
         self,
         messages: List[Dict[str, str]],
         response_format: Optional[str] = None,
+        tools: Optional[List[Dict]] = None,
+        tool_choice: str = "auto",
     ) -> str:
         """
         Generate a response based on the given messages using OpenAI.

mem0/llms/xai.py

@@ -18,13 +18,21 @@ class XAILLM(LLMBase):
         base_url = self.config.xai_base_url or os.getenv("XAI_API_BASE") or "https://api.x.ai/v1"
         self.client = OpenAI(api_key=api_key, base_url=base_url)
 
-    def generate_response(self, messages: List[Dict[str, str]], response_format=None):
+    def generate_response(
+        self,
+        messages: List[Dict[str, str]],
+        response_format=None,
+        tools: Optional[List[Dict]] = None,
+        tool_choice: str = "auto",
+    ):
         """
         Generate a response based on the given messages using XAI.
 
         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".
 
         Returns:
             str: The generated response.