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

mem0/embeddings/mock.py

@@ -0,0 +1,11 @@
+from typing import Literal, Optional
+
+from mem0.embeddings.base import EmbeddingBase
+
+
+class MockEmbeddings(EmbeddingBase):
+    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
+        """
+        Generate a mock embedding with dimension of 10.
+        """
+        return [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]

mem0/embeddings/ollama.py

@@ -0,0 +1,53 @@
+import subprocess
+import sys
+from typing import Literal, Optional
+
+from mem0.configs.embeddings.base import BaseEmbedderConfig
+from mem0.embeddings.base import EmbeddingBase
+
+try:
+    from ollama import Client
+except ImportError:
+    user_input = input("The 'ollama' library is required. Install it now? [y/N]: ")
+    if user_input.lower() == "y":
+        try:
+            subprocess.check_call([sys.executable, "-m", "pip", "install", "ollama"])
+            from ollama import Client
+        except subprocess.CalledProcessError:
+            print("Failed to install 'ollama'. Please install it manually using 'pip install ollama'.")
+            sys.exit(1)
+    else:
+        print("The required 'ollama' library is not installed.")
+        sys.exit(1)
+
+
+class OllamaEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        self.config.model = self.config.model or "nomic-embed-text"
+        self.config.embedding_dims = self.config.embedding_dims or 512
+
+        self.client = Client(host=self.config.ollama_base_url)
+        self._ensure_model_exists()
+
+    def _ensure_model_exists(self):
+        """
+        Ensure the specified model exists locally. If not, pull it from Ollama.
+        """
+        local_models = self.client.list()["models"]
+        if not any(model.get("name") == self.config.model for model in local_models):
+            self.client.pull(self.config.model)
+
+    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
+        """
+        Get the embedding for the given text using Ollama.
+
+        Args:
+            text (str): The text to embed.
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
+        Returns:
+            list: The embedding vector.
+        """
+        response = self.client.embeddings(model=self.config.model, prompt=text)
+        return response["embedding"]

mem0/embeddings/openai.py

@@ -0,0 +1,49 @@
+import os
+import warnings
+from typing import Literal, Optional
+
+from openai import OpenAI
+
+from mem0.configs.embeddings.base import BaseEmbedderConfig
+from mem0.embeddings.base import EmbeddingBase
+
+
+class OpenAIEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        self.config.model = self.config.model or "text-embedding-3-small"
+        self.config.embedding_dims = self.config.embedding_dims or 1536
+
+        api_key = self.config.api_key or os.getenv("OPENAI_API_KEY")
+        base_url = (
+            self.config.openai_base_url
+            or os.getenv("OPENAI_API_BASE")
+            or os.getenv("OPENAI_BASE_URL")
+            or "https://api.openai.com/v1"
+        )
+        if os.environ.get("OPENAI_API_BASE"):
+            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,
+            )
+
+        self.client = OpenAI(api_key=api_key, base_url=base_url)
+
+    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
+        """
+        Get the embedding for the given text using OpenAI.
+
+        Args:
+            text (str): The text to embed.
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
+        Returns:
+            list: The embedding vector.
+        """
+        text = text.replace("\n", " ")
+        return (
+            self.client.embeddings.create(input=[text], model=self.config.model, dimensions=self.config.embedding_dims)
+            .data[0]
+            .embedding
+        )

mem0/embeddings/together.py

@@ -0,0 +1,31 @@
+import os
+from typing import Literal, Optional
+
+from together import Together
+
+from mem0.configs.embeddings.base import BaseEmbedderConfig
+from mem0.embeddings.base import EmbeddingBase
+
+
+class TogetherEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        self.config.model = self.config.model or "togethercomputer/m2-bert-80M-8k-retrieval"
+        api_key = self.config.api_key or os.getenv("TOGETHER_API_KEY")
+        # TODO: check if this is correct
+        self.config.embedding_dims = self.config.embedding_dims or 768
+        self.client = Together(api_key=api_key)
+
+    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
+        """
+        Get the embedding for the given text using OpenAI.
+
+        Args:
+            text (str): The text to embed.
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
+        Returns:
+            list: The embedding vector.
+        """
+
+        return self.client.embeddings.create(model=self.config.model, input=text).data[0].embedding

mem0/embeddings/vertexai.py

@@ -0,0 +1,54 @@
+import os
+from typing import Literal, Optional
+
+from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
+
+from mem0.configs.embeddings.base import BaseEmbedderConfig
+from mem0.embeddings.base import EmbeddingBase
+
+
+class VertexAIEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        self.config.model = self.config.model or "text-embedding-004"
+        self.config.embedding_dims = self.config.embedding_dims or 256
+
+        self.embedding_types = {
+            "add": self.config.memory_add_embedding_type or "RETRIEVAL_DOCUMENT",
+            "update": self.config.memory_update_embedding_type or "RETRIEVAL_DOCUMENT",
+            "search": self.config.memory_search_embedding_type or "RETRIEVAL_QUERY",
+        }
+
+        credentials_path = self.config.vertex_credentials_json
+
+        if credentials_path:
+            os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = credentials_path
+        elif not os.getenv("GOOGLE_APPLICATION_CREDENTIALS"):
+            raise ValueError(
+                "Google application credentials JSON is not provided. Please provide a valid JSON path or set the 'GOOGLE_APPLICATION_CREDENTIALS' environment variable."
+            )
+
+        self.model = TextEmbeddingModel.from_pretrained(self.config.model)
+
+    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
+        """
+        Get the embedding for the given text using Vertex AI.
+
+        Args:
+            text (str): The text to embed.
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
+        Returns:
+            list: The embedding vector.
+        """
+        embedding_type = "SEMANTIC_SIMILARITY"
+        if memory_action is not None:
+            if memory_action not in self.embedding_types:
+                raise ValueError(f"Invalid memory action: {memory_action}")
+
+            embedding_type = self.embedding_types[memory_action]
+
+        text_input = TextEmbeddingInput(text=text, task_type=embedding_type)
+        embeddings = self.model.get_embeddings(texts=[text_input], output_dimensionality=self.config.embedding_dims)
+
+        return embeddings[0].values

mem0/graphs/configs.py

@@ -0,0 +1,96 @@
+from typing import Optional, Union
+
+from pydantic import BaseModel, Field, field_validator, model_validator
+
+from mem0.llms.configs import LlmConfig
+
+
+class Neo4jConfig(BaseModel):
+    url: Optional[str] = Field(None, description="Host address for the graph database")
+    username: Optional[str] = Field(None, description="Username for the graph database")
+    password: Optional[str] = Field(None, description="Password for the graph database")
+    database: Optional[str] = Field(None, description="Database for the graph database")
+    base_label: Optional[bool] = Field(None, description="Whether to use base node label __Entity__ for all entities")
+
+    @model_validator(mode="before")
+    def check_host_port_or_path(cls, values):
+        url, username, password = (
+            values.get("url"),
+            values.get("username"),
+            values.get("password"),
+        )
+        if not url or not username or not password:
+            raise ValueError("Please provide 'url', 'username' and 'password'.")
+        return values
+
+
+class MemgraphConfig(BaseModel):
+    url: Optional[str] = Field(None, description="Host address for the graph database")
+    username: Optional[str] = Field(None, description="Username for the graph database")
+    password: Optional[str] = Field(None, description="Password for the graph database")
+
+    @model_validator(mode="before")
+    def check_host_port_or_path(cls, values):
+        url, username, password = (
+            values.get("url"),
+            values.get("username"),
+            values.get("password"),
+        )
+        if not url or not username or not password:
+            raise ValueError("Please provide 'url', 'username' and 'password'.")
+        return values
+
+
+class NeptuneConfig(BaseModel):
+    endpoint: Optional[str] = (
+        Field(
+            None,
+            description="Endpoint to connect to a Neptune Analytics Server as neptune-graph://<graphid>",
+        ),
+    )
+    base_label: Optional[bool] = Field(None, description="Whether to use base node label __Entity__ for all entities")
+
+    @model_validator(mode="before")
+    def check_host_port_or_path(cls, values):
+        endpoint = values.get("endpoint")
+        if not endpoint:
+            raise ValueError("Please provide 'endpoint' with the format as 'neptune-graph://<graphid>'.")
+        if endpoint.startswith("neptune-db://"):
+            raise ValueError("neptune-db server is not yet supported")
+        elif endpoint.startswith("neptune-graph://"):
+            # This is a Neptune Analytics Graph
+            graph_identifier = endpoint.replace("neptune-graph://", "")
+            if not graph_identifier.startswith("g-"):
+                raise ValueError("Provide a valid 'graph_identifier'.")
+            values["graph_identifier"] = graph_identifier
+            return values
+        else:
+            raise ValueError(
+                "You must provide an endpoint to create a NeptuneServer as either neptune-db://<endpoint> or neptune-graph://<graphid>"
+            )
+
+
+class GraphStoreConfig(BaseModel):
+    provider: str = Field(
+        description="Provider of the data store (e.g., 'neo4j', 'memgraph', 'neptune')",
+        default="neo4j",
+    )
+    config: Union[Neo4jConfig, MemgraphConfig, NeptuneConfig] = Field(
+        description="Configuration for the specific data store", default=None
+    )
+    llm: Optional[LlmConfig] = Field(description="LLM configuration for querying the graph store", default=None)
+    custom_prompt: Optional[str] = Field(
+        description="Custom prompt to fetch entities from the given text", default=None
+    )
+
+    @field_validator("config")
+    def validate_config(cls, v, values):
+        provider = values.data.get("provider")
+        if provider == "neo4j":
+            return Neo4jConfig(**v.model_dump())
+        elif provider == "memgraph":
+            return MemgraphConfig(**v.model_dump())
+        elif provider == "neptune":
+            return NeptuneConfig(**v.model_dump())
+        else:
+            raise ValueError(f"Unsupported graph store provider: {provider}")