agentrun-mem0ai 0.0.11__py3-none-any.whl

agentrun_mem0/llms/gemini.py

@@ -0,0 +1,201 @@
+import os
+from typing import Dict, List, Optional
+
+try:
+    from google import genai
+    from google.genai import types
+except ImportError:
+    raise ImportError("The 'google-genai' library is required. Please install it using 'pip install google-genai'.")
+
+from agentrun_mem0.configs.llms.base import BaseLlmConfig
+from agentrun_mem0.llms.base import LLMBase
+
+
+class GeminiLLM(LLMBase):
+    def __init__(self, config: Optional[BaseLlmConfig] = None):
+        super().__init__(config)
+
+        if not self.config.model:
+            self.config.model = "gemini-2.0-flash"
+
+        api_key = self.config.api_key or os.getenv("GOOGLE_API_KEY")
+        self.client = genai.Client(api_key=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": None,
+                "tool_calls": [],
+            }
+
+            # Extract content from the first candidate
+            if response.candidates and response.candidates[0].content.parts:
+                for part in response.candidates[0].content.parts:
+                    if hasattr(part, "text") and part.text:
+                        processed_response["content"] = part.text
+                        break
+
+            # Extract function calls
+            if response.candidates and response.candidates[0].content.parts:
+                for part in response.candidates[0].content.parts:
+                    if hasattr(part, "function_call") and part.function_call:
+                        fn = part.function_call
+                        processed_response["tool_calls"].append(
+                            {
+                                "name": fn.name,
+                                "arguments": dict(fn.args) if fn.args else {},
+                            }
+                        )
+
+            return processed_response
+        else:
+            if response.candidates and response.candidates[0].content.parts:
+                for part in response.candidates[0].content.parts:
+                    if hasattr(part, "text") and part.text:
+                        return part.text
+            return ""
+
+    def _reformat_messages(self, messages: List[Dict[str, str]]):
+        """
+        Reformat messages for Gemini.
+
+        Args:
+            messages: The list of messages provided in the request.
+
+        Returns:
+            tuple: (system_instruction, contents_list)
+        """
+        system_instruction = None
+        contents = []
+
+        for message in messages:
+            if message["role"] == "system":
+                system_instruction = message["content"]
+            else:
+                content = types.Content(
+                    parts=[types.Part(text=message["content"])],
+                    role=message["role"],
+                )
+                contents.append(content)
+
+        return system_instruction, contents
+
+    def _reformat_tools(self, tools: Optional[List[Dict]]):
+        """
+        Reformat tools for Gemini.
+
+        Args:
+            tools: The list of tools provided in the request.
+
+        Returns:
+            list: The list of tools in the required format.
+        """
+
+        def remove_additional_properties(data):
+            """Recursively removes 'additionalProperties' from nested dictionaries."""
+            if isinstance(data, dict):
+                filtered_dict = {
+                    key: remove_additional_properties(value)
+                    for key, value in data.items()
+                    if not (key == "additionalProperties")
+                }
+                return filtered_dict
+            else:
+                return data
+
+        if tools:
+            function_declarations = []
+            for tool in tools:
+                func = tool["function"].copy()
+                cleaned_func = remove_additional_properties(func)
+
+                function_declaration = types.FunctionDeclaration(
+                    name=cleaned_func["name"],
+                    description=cleaned_func.get("description", ""),
+                    parameters=cleaned_func.get("parameters", {}),
+                )
+                function_declarations.append(function_declaration)
+
+            tool_obj = types.Tool(function_declarations=function_declarations)
+            return [tool_obj]
+        else:
+            return 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 Gemini.
+
+        Args:
+            messages (list): List of message dicts containing 'role' and 'content'.
+            response_format (str or object, optional): Format for 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.
+        """
+
+        # Extract system instruction and reformat messages
+        system_instruction, contents = self._reformat_messages(messages)
+
+        # Prepare generation config
+        config_params = {
+            "temperature": self.config.temperature,
+            "max_output_tokens": self.config.max_tokens,
+            "top_p": self.config.top_p,
+        }
+
+        # Add system instruction to config if present
+        if system_instruction:
+            config_params["system_instruction"] = system_instruction
+
+        if response_format is not None and response_format["type"] == "json_object":
+            config_params["response_mime_type"] = "application/json"
+            if "schema" in response_format:
+                config_params["response_schema"] = response_format["schema"]
+
+        if tools:
+            formatted_tools = self._reformat_tools(tools)
+            config_params["tools"] = formatted_tools
+
+            if tool_choice:
+                if tool_choice == "auto":
+                    mode = types.FunctionCallingConfigMode.AUTO
+                elif tool_choice == "any":
+                    mode = types.FunctionCallingConfigMode.ANY
+                else:
+                    mode = types.FunctionCallingConfigMode.NONE
+
+                tool_config = types.ToolConfig(
+                    function_calling_config=types.FunctionCallingConfig(
+                        mode=mode,
+                        allowed_function_names=(
+                            [tool["function"]["name"] for tool in tools] if tool_choice == "any" else None
+                        ),
+                    )
+                )
+                config_params["tool_config"] = tool_config
+
+        generation_config = types.GenerateContentConfig(**config_params)
+
+        response = self.client.models.generate_content(
+            model=self.config.model, contents=contents, config=generation_config
+        )
+
+        return self._parse_response(response, tools)

agentrun_mem0/llms/groq.py

@@ -0,0 +1,88 @@
+import json
+import os
+from typing import Dict, List, Optional
+
+try:
+    from groq import Groq
+except ImportError:
+    raise ImportError("The 'groq' library is required. Please install it using 'pip install groq'.")
+
+from agentrun_mem0.configs.llms.base import BaseLlmConfig
+from agentrun_mem0.llms.base import LLMBase
+from agentrun_mem0.memory.utils import extract_json
+
+
+class GroqLLM(LLMBase):
+    def __init__(self, config: Optional[BaseLlmConfig] = None):
+        super().__init__(config)
+
+        if not self.config.model:
+            self.config.model = "llama3-70b-8192"
+
+        api_key = self.config.api_key or os.getenv("GROQ_API_KEY")
+        self.client = Groq(api_key=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=None,
+        tools: Optional[List[Dict]] = None,
+        tool_choice: str = "auto",
+    ):
+        """
+        Generate a response based on the given messages using Groq.
+
+        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.
+        """
+        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 tools:
+            params["tools"] = tools
+            params["tool_choice"] = tool_choice
+
+        response = self.client.chat.completions.create(**params)
+        return self._parse_response(response, tools)

agentrun_mem0/llms/langchain.py

@@ -0,0 +1,94 @@
+from typing import Dict, List, Optional
+
+from agentrun_mem0.configs.llms.base import BaseLlmConfig
+from agentrun_mem0.llms.base import LLMBase
+
+try:
+    from langchain.chat_models.base import BaseChatModel
+    from langchain_core.messages import AIMessage
+except ImportError:
+    raise ImportError("langchain is not installed. Please install it using `pip install langchain`")
+
+
+class LangchainLLM(LLMBase):
+    def __init__(self, config: Optional[BaseLlmConfig] = None):
+        super().__init__(config)
+
+        if self.config.model is None:
+            raise ValueError("`model` parameter is required")
+
+        if not isinstance(self.config.model, BaseChatModel):
+            raise ValueError("`model` must be an instance of BaseChatModel")
+
+        self.langchain_model = self.config.model
+
+    def _parse_response(self, response: AIMessage, tools: Optional[List[Dict]]):
+        """
+        Process the response based on whether tools are used or not.
+
+        Args:
+            response: AI Message.
+            tools: The list of tools provided in the request.
+
+        Returns:
+            str or dict: The processed response.
+        """
+        if not tools:
+            return response.content
+
+        processed_response = {
+            "content": response.content,
+            "tool_calls": [],
+        }
+
+        for tool_call in response.tool_calls:
+            processed_response["tool_calls"].append(
+                {
+                    "name": tool_call["name"],
+                    "arguments": tool_call["args"],
+                }
+            )
+
+        return processed_response
+
+    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 langchain_community.
+
+        Args:
+            messages (list): List of message dicts containing 'role' and 'content'.
+            response_format (str or object, optional): Format of the response. Not used in Langchain.
+            tools (list, optional): List of tools that the model can call.
+            tool_choice (str, optional): Tool choice method.
+
+        Returns:
+            str: The generated response.
+        """
+        # Convert the messages to LangChain's tuple format
+        langchain_messages = []
+        for message in messages:
+            role = message["role"]
+            content = message["content"]
+
+            if role == "system":
+                langchain_messages.append(("system", content))
+            elif role == "user":
+                langchain_messages.append(("human", content))
+            elif role == "assistant":
+                langchain_messages.append(("ai", content))
+
+        if not langchain_messages:
+            raise ValueError("No valid messages found in the messages list")
+
+        langchain_model = self.langchain_model
+        if tools:
+            langchain_model = langchain_model.bind_tools(tools=tools, tool_choice=tool_choice)
+
+        response: AIMessage = langchain_model.invoke(langchain_messages)
+        return self._parse_response(response, tools)

agentrun_mem0/llms/litellm.py

@@ -0,0 +1,87 @@
+import json
+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'.")
+
+from agentrun_mem0.configs.llms.base import BaseLlmConfig
+from agentrun_mem0.llms.base import LLMBase
+from agentrun_mem0.memory.utils import extract_json
+
+
+class LiteLLM(LLMBase):
+    def __init__(self, config: Optional[BaseLlmConfig] = None):
+        super().__init__(config)
+
+        if not self.config.model:
+            self.config.model = "gpt-4.1-nano-2025-04-14"
+
+    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=None,
+        tools: Optional[List[Dict]] = None,
+        tool_choice: str = "auto",
+    ):
+        """
+        Generate a response based on the given messages using Litellm.
+
+        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.
+        """
+        if not litellm.supports_function_calling(self.config.model):
+            raise ValueError(f"Model '{self.config.model}' in litellm does not support function calling.")
+
+        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 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)

agentrun_mem0/llms/lmstudio.py

@@ -0,0 +1,114 @@
+import json
+from typing import Dict, List, Optional, Union
+
+from openai import OpenAI
+
+from agentrun_mem0.configs.llms.base import BaseLlmConfig
+from agentrun_mem0.configs.llms.lmstudio import LMStudioConfig
+from agentrun_mem0.llms.base import LLMBase
+from agentrun_mem0.memory.utils import extract_json
+
+
+class LMStudioLLM(LLMBase):
+    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 = (
+            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)
+
+    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=None,
+        tools: Optional[List[Dict]] = None,
+        tool_choice: str = "auto",
+        **kwargs,
+    ):
+        """
+        Generate a response based on the given messages using LM Studio.
+
+        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 LM Studio-specific parameters.
+
+        Returns:
+            str: The generated response.
+        """
+        params = self._get_supported_params(messages=messages, **kwargs)
+        params.update(
+            {
+                "model": self.config.model,
+                "messages": messages,
+            }
+        )
+
+        if self.config.lmstudio_response_format:
+            params["response_format"] = self.config.lmstudio_response_format
+        elif response_format:
+            params["response_format"] = response_format
+        else:
+            params["response_format"] = {"type": "json_object"}
+
+        if tools:
+            params["tools"] = tools
+            params["tool_choice"] = tool_choice
+
+        response = self.client.chat.completions.create(**params)
+        return self._parse_response(response, tools)