mem0ai-azure-mysql 0.1.115__py3-none-any.whl

mem0/llms/aws_bedrock.py

@@ -0,0 +1,270 @@
+import json
+import os
+import re
+from typing import Any, Dict, List, Optional
+
+try:
+    import boto3
+except ImportError:
+    raise ImportError("The 'boto3' library is required. Please install it using 'pip install boto3'.")
+
+from mem0.configs.llms.base import BaseLlmConfig
+from mem0.llms.base import LLMBase
+
+PROVIDERS = ["ai21", "amazon", "anthropic", "cohere", "meta", "mistral", "stability", "writer"]
+
+
+def extract_provider(model: str) -> str:
+    for provider in PROVIDERS:
+        if re.search(rf"\b{re.escape(provider)}\b", model):
+            return provider
+    raise ValueError(f"Unknown provider in model: {model}")
+
+
+class AWSBedrockLLM(LLMBase):
+    def __init__(self, config: Optional[BaseLlmConfig] = None):
+        super().__init__(config)
+
+        if not self.config.model:
+            self.config.model = "anthropic.claude-3-5-sonnet-20240620-v1:0"
+
+        # Get AWS config from environment variables or use defaults
+        aws_access_key = os.environ.get("AWS_ACCESS_KEY_ID", "")
+        aws_secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY", "")
+        aws_region = os.environ.get("AWS_REGION", "us-west-2")
+
+        # Check if AWS config is provided in the config
+        if hasattr(self.config, "aws_access_key_id"):
+            aws_access_key = self.config.aws_access_key_id
+        if hasattr(self.config, "aws_secret_access_key"):
+            aws_secret_key = self.config.aws_secret_access_key
+        if hasattr(self.config, "aws_region"):
+            aws_region = self.config.aws_region
+
+        self.client = boto3.client(
+            "bedrock-runtime",
+            region_name=aws_region,
+            aws_access_key_id=aws_access_key if aws_access_key else None,
+            aws_secret_access_key=aws_secret_key if aws_secret_key else None,
+        )
+
+        self.model_kwargs = {
+            "temperature": self.config.temperature,
+            "max_tokens_to_sample": self.config.max_tokens,
+            "top_p": self.config.top_p,
+        }
+
+    def _format_messages(self, messages: List[Dict[str, str]]) -> str:
+        """
+        Formats a list of messages into the required prompt structure for the model.
+
+        Args:
+            messages (List[Dict[str, str]]): A list of dictionaries where each dictionary represents a message.
+                                            Each dictionary contains 'role' and 'content' keys.
+
+        Returns:
+            str: A formatted string combining all messages, structured with roles capitalized and separated by newlines.
+        """
+
+        formatted_messages = []
+        for message in messages:
+            role = message["role"].capitalize()
+            content = message["content"]
+            formatted_messages.append(f"\n\n{role}: {content}")
+
+        return "\n\nHuman: " + "".join(formatted_messages) + "\n\nAssistant:"
+
+    def _parse_response(self, response, tools) -> str:
+        """
+        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 = {"tool_calls": []}
+
+            if response["output"]["message"]["content"]:
+                for item in response["output"]["message"]["content"]:
+                    if "toolUse" in item:
+                        processed_response["tool_calls"].append(
+                            {
+                                "name": item["toolUse"]["name"],
+                                "arguments": item["toolUse"]["input"],
+                            }
+                        )
+
+            return processed_response
+
+        response_body = response.get("body").read().decode()
+        response_json = json.loads(response_body)
+        return response_json.get("content", [{"text": ""}])[0].get("text", "")
+
+    def _prepare_input(
+        self,
+        provider: str,
+        model: str,
+        prompt: str,
+        model_kwargs: Optional[Dict[str, Any]] = {},
+    ) -> Dict[str, Any]:
+        """
+        Prepares the input dictionary for the specified provider's model by mapping and renaming
+        keys in the input based on the provider's requirements.
+
+        Args:
+            provider (str): The name of the service provider (e.g., "meta", "ai21", "mistral", "cohere", "amazon").
+            model (str): The name or identifier of the model being used.
+            prompt (str): The text prompt to be processed by the model.
+            model_kwargs (Dict[str, Any]): Additional keyword arguments specific to the model's requirements.
+
+        Returns:
+            Dict[str, Any]: The prepared input dictionary with the correct keys and values for the specified provider.
+        """
+
+        input_body = {"prompt": prompt, **model_kwargs}
+
+        provider_mappings = {
+            "meta": {"max_tokens_to_sample": "max_gen_len"},
+            "ai21": {"max_tokens_to_sample": "maxTokens", "top_p": "topP"},
+            "mistral": {"max_tokens_to_sample": "max_tokens"},
+            "cohere": {"max_tokens_to_sample": "max_tokens", "top_p": "p"},
+        }
+
+        if provider in provider_mappings:
+            for old_key, new_key in provider_mappings[provider].items():
+                if old_key in input_body:
+                    input_body[new_key] = input_body.pop(old_key)
+
+        if provider == "cohere" and "cohere.command-r" in model:
+            input_body["message"] = input_body.pop("prompt")
+
+        if provider == "amazon":
+            input_body = {
+                "inputText": prompt,
+                "textGenerationConfig": {
+                    "maxTokenCount": self.model_kwargs["max_tokens_to_sample"]
+                    or self.model_kwargs["max_tokens"]
+                    or 5000,
+                    "topP": self.model_kwargs["top_p"] or 0.9,
+                    "temperature": self.model_kwargs["temperature"] or 0.1,
+                },
+            }
+            input_body["textGenerationConfig"] = {
+                k: v for k, v in input_body["textGenerationConfig"].items() if v is not None
+            }
+
+        return input_body
+
+    def _convert_tool_format(self, original_tools):
+        """
+        Converts a list of tools from their original format to a new standardized format.
+
+        Args:
+            original_tools (list): A list of dictionaries representing the original tools, each containing a 'type' key and corresponding details.
+
+        Returns:
+            list: A list of dictionaries representing the tools in the new standardized format.
+        """
+        new_tools = []
+
+        for tool in original_tools:
+            if tool["type"] == "function":
+                function = tool["function"]
+                new_tool = {
+                    "toolSpec": {
+                        "name": function["name"],
+                        "description": function["description"],
+                        "inputSchema": {
+                            "json": {
+                                "type": "object",
+                                "properties": {},
+                                "required": function["parameters"].get("required", []),
+                            }
+                        },
+                    }
+                }
+
+                for prop, details in function["parameters"].get("properties", {}).items():
+                    new_tool["toolSpec"]["inputSchema"]["json"]["properties"][prop] = details
+
+                new_tools.append(new_tool)
+
+        return new_tools
+
+    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 AWS Bedrock.
+
+        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.
+        """
+
+        if tools:
+            # Use converse method when tools are provided
+            messages = [
+                {
+                    "role": "user",
+                    "content": [{"text": message["content"]} for message in messages],
+                }
+            ]
+            inference_config = {
+                "temperature": self.model_kwargs["temperature"],
+                "maxTokens": self.model_kwargs["max_tokens_to_sample"],
+                "topP": self.model_kwargs["top_p"],
+            }
+            tools_config = {"tools": self._convert_tool_format(tools)}
+
+            response = self.client.converse(
+                modelId=self.config.model,
+                messages=messages,
+                inferenceConfig=inference_config,
+                toolConfig=tools_config,
+            )
+        else:
+            # Use invoke_model method when no tools are provided
+            prompt = self._format_messages(messages)
+            provider = extract_provider(self.config.model)
+            input_body = self._prepare_input(provider, self.config.model, prompt, model_kwargs=self.model_kwargs)
+            body = json.dumps(input_body)
+
+            if provider == "anthropic" or provider == "deepseek":
+                input_body = {
+                    "messages": [{"role": "user", "content": [{"type": "text", "text": prompt}]}],
+                    "max_tokens": self.model_kwargs["max_tokens_to_sample"] or self.model_kwargs["max_tokens"] or 5000,
+                    "temperature": self.model_kwargs["temperature"] or 0.1,
+                    "top_p": self.model_kwargs["top_p"] or 0.9,
+                    "anthropic_version": "bedrock-2023-05-31",
+                }
+
+                body = json.dumps(input_body)
+
+                response = self.client.invoke_model(
+                    body=body,
+                    modelId=self.config.model,
+                    accept="application/json",
+                    contentType="application/json",
+                )
+            else:
+                response = self.client.invoke_model(
+                    body=body,
+                    modelId=self.config.model,
+                    accept="application/json",
+                    contentType="application/json",
+                )
+
+        return self._parse_response(response, tools)

mem0/llms/azure_openai.py

@@ -0,0 +1,114 @@
+import json
+import os
+from typing import Dict, List, Optional
+
+from openai import AzureOpenAI
+from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+
+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):
+        super().__init__(config)
+
+        # Model name should match the custom deployment name chosen for it.
+        if not self.config.model:
+            self.config.model = "gpt-4o"
+
+        azure_deployment = self.config.azure_kwargs.azure_deployment or os.getenv("LLM_AZURE_DEPLOYMENT")
+        azure_endpoint = self.config.azure_kwargs.azure_endpoint or os.getenv("LLM_AZURE_ENDPOINT")
+        api_version = self.config.azure_kwargs.api_version or os.getenv("LLM_AZURE_API_VERSION")
+        default_headers = self.config.azure_kwargs.default_headers
+
+        credential = DefaultAzureCredential()
+        token_provider = get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
+        self.client = AzureOpenAI(
+            azure_deployment=azure_deployment,
+            azure_endpoint=azure_endpoint,
+            api_version=api_version,
+            azure_ad_token_provider=token_provider,
+            http_client=self.config.http_client,
+            default_headers=default_headers,
+        )
+
+    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 Azure OpenAI.
+
+        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.
+        """
+
+        user_prompt = messages[-1]["content"]
+
+        user_prompt = user_prompt.replace("assistant", "ai")
+
+        messages[-1]["content"] = user_prompt
+
+        common_params = {
+            "model": self.config.model,
+            "messages": messages,
+        }
+
+        if self.config.model in {"o3-mini", "o1-preview", "o1"}:
+            params = common_params
+        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
+            params["tools"] = tools
+            params["tool_choice"] = tool_choice
+
+        response = self.client.chat.completions.create(**params)
+        return self._parse_response(response, tools)

mem0/llms/azure_openai_structured.py

@@ -0,0 +1,76 @@
+import os
+from typing import Dict, List, Optional
+
+from openai import AzureOpenAI
+
+from mem0.configs.llms.base import BaseLlmConfig
+from mem0.llms.base import LLMBase
+
+
+class AzureOpenAIStructuredLLM(LLMBase):
+    def __init__(self, config: Optional[BaseLlmConfig] = None):
+        super().__init__(config)
+
+        # Model name should match the custom deployment name chosen for it.
+        if not self.config.model:
+            self.config.model = "gpt-4o-2024-08-06"
+
+        api_key = os.getenv("LLM_AZURE_OPENAI_API_KEY") or self.config.azure_kwargs.api_key
+        azure_deployment = os.getenv("LLM_AZURE_DEPLOYMENT") or self.config.azure_kwargs.azure_deployment
+        azure_endpoint = os.getenv("LLM_AZURE_ENDPOINT") or self.config.azure_kwargs.azure_endpoint
+        api_version = os.getenv("LLM_AZURE_API_VERSION") or self.config.azure_kwargs.api_version
+        default_headers = self.config.azure_kwargs.default_headers
+
+        # Can display a warning if API version is of model and api-version
+        self.client = AzureOpenAI(
+            azure_deployment=azure_deployment,
+            azure_endpoint=azure_endpoint,
+            api_version=api_version,
+            api_key=api_key,
+            http_client=self.config.http_client,
+            default_headers=default_headers,
+        )
+
+    def generate_response(
+        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.
+
+        Args:
+            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.
+        """
+
+        user_prompt = messages[-1]["content"]
+
+        user_prompt = user_prompt.replace("assistant", "ai")
+
+        messages[-1]["content"] = user_prompt
+
+        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
+
+        if tools:
+            params["tools"] = tools
+            params["tool_choice"] = tool_choice
+
+        response = self.client.chat.completions.create(**params)
+        return self._parse_response(response, tools)

mem0/llms/base.py

@@ -0,0 +1,32 @@
+from abc import ABC, abstractmethod
+from typing import Dict, List, Optional
+
+from mem0.configs.llms.base import BaseLlmConfig
+
+
+class LLMBase(ABC):
+    def __init__(self, config: Optional[BaseLlmConfig] = None):
+        """Initialize a base LLM class
+
+        :param config: LLM configuration option class, defaults to None
+        :type config: Optional[BaseLlmConfig], optional
+        """
+        if config is None:
+            self.config = BaseLlmConfig()
+        else:
+            self.config = config
+
+    @abstractmethod
+    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.
+        """
+        pass

mem0/llms/configs.py

@@ -0,0 +1,34 @@
+from typing import Optional
+
+from pydantic import BaseModel, Field, field_validator
+
+
+class LlmConfig(BaseModel):
+    provider: str = Field(description="Provider of the LLM (e.g., 'ollama', 'openai')", default="openai")
+    config: Optional[dict] = Field(description="Configuration for the specific LLM", default={})
+
+    @field_validator("config")
+    def validate_config(cls, v, values):
+        provider = values.data.get("provider")
+        if provider in (
+            "openai",
+            "ollama",
+            "anthropic",
+            "groq",
+            "together",
+            "aws_bedrock",
+            "litellm",
+            "azure_openai",
+            "openai_structured",
+            "azure_openai_structured",
+            "gemini",
+            "deepseek",
+            "xai",
+            "sarvam",
+            "lmstudio",
+            "vllm",
+            "langchain",
+        ):
+            return v
+        else:
+            raise ValueError(f"Unsupported LLM provider: {provider}")

mem0/llms/deepseek.py

@@ -0,0 +1,85 @@
+import json
+import os
+from typing import Dict, List, Optional
+
+from openai import OpenAI
+
+from mem0.configs.llms.base import BaseLlmConfig
+from mem0.llms.base import LLMBase
+from mem0.memory.utils import extract_json
+
+
+class DeepSeekLLM(LLMBase):
+    def __init__(self, config: Optional[BaseLlmConfig] = None):
+        super().__init__(config)
+
+        if not self.config.model:
+            self.config.model = "deepseek-chat"
+
+        api_key = self.config.api_key or os.getenv("DEEPSEEK_API_KEY")
+        base_url = self.config.deepseek_base_url or os.getenv("DEEPSEEK_API_BASE") or "https://api.deepseek.com"
+        self.client = OpenAI(api_key=api_key, base_url=base_url)
+
+    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 DeepSeek.
+
+        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 tools:
+            params["tools"] = tools
+            params["tool_choice"] = tool_choice
+
+        response = self.client.chat.completions.create(**params)
+        return self._parse_response(response, tools)